From 8c5de70db10895c0a8afc07223cb33703d1c0920 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Tue, 18 Aug 2026 16:20:27 +0300 Subject: [PATCH 01/25] fix: stop demo query streams from dying mid-flight (research#86) Three consecutive demo queries failed live on 2026-07-29 with "Stream error: network error". The backend never errored; the response body went silent for the whole SQL-generation phase and the connection was severed mid-flight. Streaming (the incident): - Add `with_keepalive`, wrapping the serialized stream so a bare delimiter is emitted every 10s while the pipeline produces nothing. A bare delimiter splits into an empty part, which every existing client parser already skips, so this needs no protocol or client change. Applied to all four streaming endpoints (query, confirm, refresh, connect-database). - Set `Cache-Control: no-cache, no-transform` and `X-Accel-Buffering: no` to discourage intermediaries from buffering the body. The media type stays `application/json`: the wire format is delimited JSON, not SSE, so declaring `text/event-stream` would misdescribe it. Migrating to real SSE is a follow-up. - Move every synchronous LLM call off the event loop via `asyncio.to_thread`: `get_analysis`, `heal_and_execute`, the follow-up agent and both `format_ai_response` calls. `RelevancyAgent.get_answer` was `async def` but called `run_completion` synchronously, so its `create_task` concurrency with table-finding was illusory and it blocked the loop too. This is why the failures clustered across users rather than hitting one request. Instrumentation (why it stayed undiagnosable): - `run_completion` now applies `Config.LLM_TIMEOUT` (default 90s), passed to litellm so it aborts the HTTP request rather than hanging forever, and logs every call's duration with a caller label. Calls over `LLM_SLOW_CALL_THRESHOLD` (default 20s) log at WARNING. The analysis agent had zero instrumentation, so the original slowness left no trace at all. - Route `HealerAgent` through `run_completion` so it inherits both; it called `litellm.completion` directly and had no timeout. UI: - The `sqlQuery !== undefined` render guard was always true, since `sqlQuery` is initialized to `""`. Failed runs painted an empty "Query Analysis" card, which made the screenshots misleading. Guard on truthiness. Memory (present in the same logs, unrelated to the failure): - Default `AZURE_API_VERSION` to `2025-03-01-preview`. Graphiti's client uses the Azure Responses API, which rejects older versions with HTTP 400, so every episode write was failing. - `len(history[1])` threw on the first message of a session, where the client sends no result array. Use a falsy check. - Log the previously silent `except` in `update_user_information`, which hid the failure on that path entirely. Tests: 6 new unit tests for the keepalive wrapper covering pass-through, silent-gap emission, client-parser compatibility, exception propagation and teardown on client disconnect. Verified at the wire level against uvicorn: keepalive frames arrive every ~0.4s through a 2s silent gap. Refs: research#86, incident 2026-07-29 Co-Authored-By: Claude Opus 5 (1M context) --- .env.example | 12 ++- api/agents/analysis_agent.py | 3 +- api/agents/follow_up_agent.py | 3 +- api/agents/healer_agent.py | 12 +-- api/agents/relevancy_agent.py | 11 ++- api/agents/response_formatter_agent.py | 1 + api/agents/utils.py | 33 ++++++- api/config.py | 12 +++ api/core/text2sql.py | 22 +++-- api/memory/graphiti_tool.py | 12 ++- api/routes/database.py | 7 +- api/routes/graphs.py | 23 +++-- api/routes/streaming.py | 62 +++++++++++++ app/src/components/chat/ChatInterface.tsx | 7 +- tests/test_stream_keepalive.py | 103 ++++++++++++++++++++++ 15 files changed, 290 insertions(+), 33 deletions(-) create mode 100644 api/routes/streaming.py create mode 100644 tests/test_stream_keepalive.py diff --git a/.env.example b/.env.example index eeecd753..6ef6ec98 100644 --- a/.env.example +++ b/.env.example @@ -75,6 +75,14 @@ FALKORDB_URL=redis://localhost:6379/0 # REQUIRED - change to your FalkorDB URL # COMPLETION_MODEL=openai/gpt-4.1 # EMBEDDING_MODEL=openai/text-embedding-ada-002 +# Wall-clock ceiling for a single agent LLM call, in seconds (default 90). +# Passed to litellm, which aborts the HTTP request — a hung provider then +# surfaces as an error instead of stalling the response stream. +# LLM_TIMEOUT=90 +# +# Calls slower than this are logged at WARNING (default 20). +# LLM_SLOW_CALL_THRESHOLD=20 + # OpenAI - uses openai/gpt-4.1 and openai/text-embedding-ada-002 # OPENAI_API_KEY=your_openai_api_key @@ -100,7 +108,9 @@ FALKORDB_URL=redis://localhost:6379/0 # REQUIRED - change to your FalkorDB URL # Azure OpenAI (default fallback) - uses azure/gpt-4.1 and azure/text-embedding-ada-002 # AZURE_API_KEY=your_azure_api_key # AZURE_API_BASE=https://your-resource.openai.azure.com/ -# AZURE_API_VERSION=2023-05-15 +# Must be 2025-03-01-preview or later — Graphiti memory writes use the +# Azure Responses API, which rejects older api-versions with HTTP 400. +# AZURE_API_VERSION=2025-03-01-preview # ----------------------------- # OAuth configuration (optional — uncomment to enable login flows) diff --git a/api/agents/analysis_agent.py b/api/agents/analysis_agent.py index 9ff49a09..a7e57d21 100644 --- a/api/agents/analysis_agent.py +++ b/api/agents/analysis_agent.py @@ -38,7 +38,8 @@ def get_analysis( # pylint: disable=too-many-arguments, too-many-positional-arg self.messages.append({"role": "user", "content": prompt}) response = run_completion( - self.messages, self.custom_model, self.custom_api_key, temperature=0 + self.messages, self.custom_model, self.custom_api_key, + label="analysis", temperature=0, ) analysis = parse_response(response) if isinstance(analysis["ambiguities"], list): diff --git a/api/agents/follow_up_agent.py b/api/agents/follow_up_agent.py index b39661ec..0e0de256 100644 --- a/api/agents/follow_up_agent.py +++ b/api/agents/follow_up_agent.py @@ -70,7 +70,8 @@ def generate_follow_up_question( try: response = run_completion( [{"role": "user", "content": prompt}], - self.custom_model, self.custom_api_key, temperature=0.9 + self.custom_model, self.custom_api_key, + label="followup", temperature=0.9, ) return response.strip() diff --git a/api/agents/healer_agent.py b/api/agents/healer_agent.py index e0ab66a6..fcd4265e 100644 --- a/api/agents/healer_agent.py +++ b/api/agents/healer_agent.py @@ -10,9 +10,7 @@ import re from typing import Dict, Callable, Any -from litellm import completion -from api.config import Config -from .utils import parse_response +from .utils import parse_response, run_completion class HealerAgent: @@ -224,14 +222,12 @@ def heal_and_execute( # pylint: disable=too-many-locals for attempt in range(self.max_healing_attempts): # Call LLM - response = completion( - model=Config.COMPLETION_MODEL, - messages=self.messages, + content = run_completion( + self.messages, + label=f"healer.attempt{attempt + 1}", temperature=0.1, max_tokens=2000 ) - - content = response.choices[0].message.content self.messages.append({"role": "assistant", "content": content}) # Parse response diff --git a/api/agents/relevancy_agent.py b/api/agents/relevancy_agent.py index 84b0328d..8352c47a 100644 --- a/api/agents/relevancy_agent.py +++ b/api/agents/relevancy_agent.py @@ -1,5 +1,6 @@ """Relevancy agent for determining relevancy of queries to database schema.""" +import asyncio import json from .utils import BaseAgent, parse_response, run_completion @@ -82,8 +83,14 @@ async def get_answer(self, user_question: str, database_desc: dict) -> dict: } ) - answer = run_completion( - self.messages, self.custom_model, self.custom_api_key, temperature=0 + # ``run_completion`` is synchronous. Awaiting it off-loop matters even + # though this method is already ``async``: the caller runs it as a task + # alongside table-finding, and a blocking call here would stall that + # task — and every other request — rather than overlap with it. + answer = await asyncio.to_thread( + run_completion, + self.messages, self.custom_model, self.custom_api_key, + label="relevancy", temperature=0, ) self.messages.append({"role": "assistant", "content": answer}) return parse_response(answer) diff --git a/api/agents/response_formatter_agent.py b/api/agents/response_formatter_agent.py index 9e9dfe87..40306278 100644 --- a/api/agents/response_formatter_agent.py +++ b/api/agents/response_formatter_agent.py @@ -77,6 +77,7 @@ def format_response(self, user_query: str, sql_query: str, response = run_completion( messages, self.custom_model, self.custom_api_key, + label="formatter", temperature=0.3 # Slightly higher temperature for more natural responses ) return response.strip() diff --git a/api/agents/utils.py b/api/agents/utils.py index bc28c99f..68954eae 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -1,6 +1,8 @@ """Utility functions for agents.""" import json +import logging +import time from typing import Any, Dict, List from litellm import completion @@ -8,22 +10,49 @@ def run_completion(messages: List[Dict[str, str]], custom_model: str = None, - custom_api_key: str = None, **kwargs) -> str: + custom_api_key: str = None, *, label: str = "llm", + **kwargs) -> str: """Run an LLM completion with optional custom model/key overrides. + Applies ``Config.LLM_TIMEOUT`` unless the caller passes an explicit + ``timeout``, and logs the call duration. Both exist because the 2026-07-29 + demo failure was an LLM call that stalled with no timeout and left no + trace of how long it ran. ``label`` names the caller in those log lines + and is not forwarded to the provider. + Returns the content string from the first choice. """ completion_args = { "model": custom_model if custom_model else Config.COMPLETION_MODEL, "messages": messages, "top_p": 1, + "timeout": Config.LLM_TIMEOUT, **kwargs, } if custom_api_key: completion_args["api_key"] = custom_api_key - result = completion(**completion_args) + started = time.monotonic() + try: + result = completion(**completion_args) + except Exception: + logging.warning( + "llm_call label=%s model=%s duration=%.2fs outcome=error", + label, completion_args["model"], time.monotonic() - started, + ) + raise + elapsed = time.monotonic() - started + logging.info( + "llm_call label=%s model=%s duration=%.2fs outcome=ok", + label, completion_args["model"], elapsed, + ) + if elapsed >= Config.LLM_SLOW_CALL_THRESHOLD: + logging.warning( + "llm_call label=%s model=%s duration=%.2fs exceeded slow-call " + "threshold of %.0fs", label, completion_args["model"], elapsed, + Config.LLM_SLOW_CALL_THRESHOLD, + ) return result.choices[0].message.content diff --git a/api/config.py b/api/config.py index dce8a34c..a29c17f0 100644 --- a/api/config.py +++ b/api/config.py @@ -133,6 +133,18 @@ class Config: COMPLETION_MODEL = _user_completion or "azure/gpt-4.1" EMBEDDING_MODEL_NAME = _user_embedding or "azure/text-embedding-ada-002" + # Wall-clock ceiling for a single agent LLM call, in seconds. Passed + # through to litellm, which aborts the underlying HTTP request — so a + # hung provider surfaces as a clean error instead of stalling the + # response stream indefinitely (incident 2026-07-29). + LLM_TIMEOUT: float = float(os.getenv("LLM_TIMEOUT", "90")) # pylint: disable=invalid-name + + # A call slower than this is logged at WARNING. Normal analysis calls + # completed in ~6s during the incident window, so this flags outliers + # well before they reach the timeout. + # pylint: disable-next=invalid-name + LLM_SLOW_CALL_THRESHOLD: float = float(os.getenv("LLM_SLOW_CALL_THRESHOLD", "20")) + DB_MAX_DISTINCT: int = 100 # pylint: disable=invalid-name DB_UNIQUENESS_THRESHOLD: float = 0.5 # pylint: disable=invalid-name SHORT_MEMORY_LENGTH = 5 # Maximum number of questions to keep in short-term memory diff --git a/api/core/text2sql.py b/api/core/text2sql.py index 90cd9cdc..5f9fb73d 100644 --- a/api/core/text2sql.py +++ b/api/core/text2sql.py @@ -407,7 +407,12 @@ async def run_query( # pylint: disable=too-many-locals,too-many-branches,too-ma agent_an = AnalysisAgent( queries_history, result_history, custom_api_key, custom_model, ) - answer_an = agent_an.get_analysis( + # ``get_analysis`` is a synchronous LLM call. Running it directly here + # would block the event loop for its full duration, stalling every other + # in-flight request and preventing any stream from flushing bytes + # (incident 2026-07-29). Off-loop via a worker thread. + answer_an = await asyncio.to_thread( + agent_an.get_analysis, queries_history[-1], tables, db_description, instructions, memory_context, db_type, user_rules_spec, ) @@ -427,7 +432,8 @@ async def run_query( # pylint: disable=too-many-locals,too-many-branches,too-ma follow_up_agent = FollowUpAgent( queries_history, result_history, custom_api_key, custom_model, ) - follow_up = follow_up_agent.generate_follow_up_question( + follow_up = await asyncio.to_thread( + follow_up_agent.generate_follow_up_question, user_question=queries_history[-1], analysis_result=answer_an, ) @@ -536,7 +542,11 @@ def _run_sql(sql: str): ) return loader_class.execute_sql_query(sql, db_url) - healing_result = healer.heal_and_execute( + # Same reasoning as ``get_analysis`` above, and worse here: this + # chains up to ``max_healing_attempts`` sequential LLM calls plus + # SQL execution, all synchronous. + healing_result = await asyncio.to_thread( + healer.heal_and_execute, initial_sql=sql_query, initial_error=str(exec_error), execute_sql_func=_run_sql, @@ -593,7 +603,8 @@ def _run_sql(sql: str): "message": f"Step {step_num}: Generating user-friendly response", } - user_readable_response = format_ai_response( + user_readable_response = await asyncio.to_thread( + format_ai_response, queries_history=queries_history, result_history=result_history, sql_query=sql_query, @@ -755,7 +766,8 @@ async def run_confirmed( # pylint: disable=too-many-locals,too-many-branches,to yield {"type": "reasoning_step", "message": f"Step {step_num}: Generating user-friendly response"} - user_readable_response = format_ai_response( + user_readable_response = await asyncio.to_thread( + format_ai_response, queries_history=queries_history or [question], result_history=None, sql_query=sql_query, diff --git a/api/memory/graphiti_tool.py b/api/memory/graphiti_tool.py index 1a052c57..96665c6d 100644 --- a/api/memory/graphiti_tool.py +++ b/api/memory/graphiti_tool.py @@ -253,7 +253,7 @@ async def update_user_information(self, conversation: Dict[str, Any], history: T """ try: - if len(history[1]) == 0: + if not history[1]: messages = [{"role": "user", "content": prompt}] else: messages = [] @@ -277,6 +277,9 @@ async def update_user_information(self, conversation: Dict[str, Any], history: T await driver.execute_query(query, user_id=self.user_id, summary=content) return True except Exception as e: + # Previously swallowed silently, which hid a recurring failure on + # this path entirely (incident 2026-07-29). + logging.error("Error updating user information: %s", e) return False async def add_new_memory(self, conversation: Dict[str, Any], history: Tuple[List[str], List[str]]) -> bool: @@ -733,7 +736,7 @@ async def summarize_conversation(self, conversation: Dict[str, Any], history: Li try: - if len(history[1]) == 0: + if not history[1]: messages = [{"role": "user", "content": prompt}] else: messages = [] @@ -769,7 +772,10 @@ def __init__(self): self.api_key = os.getenv('AZURE_API_KEY') self.endpoint = os.getenv('AZURE_API_BASE') - self.api_version = os.getenv('AZURE_API_VERSION', '2024-02-01') + # Graphiti's OpenAI client uses the Responses API, which requires + # api-version 2025-03-01-preview or later. Older values make every + # episode write fail with HTTP 400 (incident 2026-07-29). + self.api_version = os.getenv('AZURE_API_VERSION', '2025-03-01-preview') self.model_choice = "gpt-4.1" # Use the model name directly # Extract just the model name without provider prefix for Graphiti diff --git a/api/routes/database.py b/api/routes/database.py index e3287541..e88e1ae3 100644 --- a/api/routes/database.py +++ b/api/routes/database.py @@ -5,6 +5,7 @@ from api.auth.user_management import token_required from api.core.schema_loader import load_database +from api.routes.streaming import STREAM_HEADERS, with_keepalive from api.routes.tokens import UNAUTHORIZED_RESPONSE database_router = APIRouter(tags=["Database Connection"]) @@ -30,4 +31,8 @@ async def connect_database(request: Request, db_request: DatabaseConnectionReque Requires authentication. """ generator = await load_database(db_request.url, request.state.user_id) - return StreamingResponse(generator, media_type="application/json") + return StreamingResponse( + with_keepalive(generator), + media_type="application/json", + headers=STREAM_HEADERS, + ) diff --git a/api/routes/graphs.py b/api/routes/graphs.py index 38003faf..96a80a03 100644 --- a/api/routes/graphs.py +++ b/api/routes/graphs.py @@ -31,6 +31,7 @@ from api.auth.user_management import token_required from api.routes.tokens import UNAUTHORIZED_RESPONSE from api.routes.usage_tracking import record_query_usage_background +from api.routes.streaming import STREAM_HEADERS, with_keepalive graphs_router = APIRouter(tags=["Graphs & Databases"]) @@ -205,14 +206,14 @@ async def stream(): question = chat_data.chat[-1] query_id = str(uuid.uuid4()) try: - async for chunk in _serialize_pipeline( + async for chunk in with_keepalive(_serialize_pipeline( run_query(request.state.user_id, graph_id, chat_data), user_id=request.state.user_id, namespaced=namespaced, question=question, query_id=query_id, endpoint=request.url.path, - ): + )): yield chunk except Exception: # pylint: disable=broad-exception-caught # Don't leak stack traces (CodeQL: information exposure through @@ -235,7 +236,9 @@ async def stream(): "message": "Internal error while processing query", }) + MESSAGE_DELIMITER - return StreamingResponse(stream(), media_type="application/json") + return StreamingResponse( + stream(), media_type="application/json", headers=STREAM_HEADERS, + ) @graphs_router.post("/{graph_id}/confirm", responses={401: UNAUTHORIZED_RESPONSE}) @@ -268,14 +271,14 @@ async def stream(): question = str(confirm_data.chat[-1]) if confirm_data.chat else "" query_id = str(uuid.uuid4()) try: - async for chunk in _serialize_pipeline( + async for chunk in with_keepalive(_serialize_pipeline( run_confirmed(request.state.user_id, graph_id, confirm_data), user_id=request.state.user_id, namespaced=namespaced, question=question, query_id=query_id, endpoint=request.url.path, - ): + )): yield chunk except Exception: # pylint: disable=broad-exception-caught # See note on the query endpoint above (CodeQL). @@ -296,7 +299,9 @@ async def stream(): "message": "Internal error while processing confirmation", }) + MESSAGE_DELIMITER - return StreamingResponse(stream(), media_type="application/json") + return StreamingResponse( + stream(), media_type="application/json", headers=STREAM_HEADERS, + ) @graphs_router.post("/{graph_id}/refresh", responses={401: UNAUTHORIZED_RESPONSE}) @@ -310,7 +315,11 @@ async def refresh_graph_schema(request: Request, graph_id: str): """ try: generator = await refresh_database_schema(request.state.user_id, graph_id) - return StreamingResponse(generator, media_type="application/json") + return StreamingResponse( + with_keepalive(generator), + media_type="application/json", + headers=STREAM_HEADERS, + ) except (InternalError, InvalidArgumentError) as e: # Log detailed error internally, send generic message to user if isinstance(e, InternalError): diff --git a/api/routes/streaming.py b/api/routes/streaming.py new file mode 100644 index 00000000..50c6be2c --- /dev/null +++ b/api/routes/streaming.py @@ -0,0 +1,62 @@ +"""Keepalive support for the delimited streaming responses. + +The query pipeline emits no events for the whole SQL-generation phase. An HTTP +body that goes silent for that long invites proxy buffering and idle-timeout +disconnects, which is what broke a live demo on 2026-07-29: the stream was +severed mid-body and the browser surfaced it as ``Stream error: network +error``. + +This lives in the route layer rather than ``api/core`` because it is a +transport concern, and so it ships with the hosted app rather than the SDK. +""" + +import asyncio + +from api.core.pipeline import MESSAGE_DELIMITER + +# Interval between keepalive bytes on an otherwise silent stream. Comfortably +# under the ~60s idle timeout common to proxies and PaaS edges. +STREAM_KEEPALIVE_INTERVAL = 10.0 + +# Discourage intermediaries from buffering the streamed body. +# ``X-Accel-Buffering`` is honoured by nginx and several PaaS edges. +STREAM_HEADERS = { + "Cache-Control": "no-cache, no-transform", + "X-Accel-Buffering": "no", +} + + +async def with_keepalive(chunks, interval: float = STREAM_KEEPALIVE_INTERVAL): + """Emit a bare delimiter while *chunks* produces nothing. + + Wrap the already-serialized stream, so one call covers a whole endpoint + including silent gaps introduced later. A bare delimiter splits into an + empty part on the client, which every consumer's parser already skips, so + this needs no protocol change and no client change. + """ + iterator = aiter(chunks) + try: + while True: + pending = asyncio.ensure_future(anext(iterator)) + try: + while True: + done, _ = await asyncio.wait({pending}, timeout=interval) + if done: + break + yield MESSAGE_DELIMITER + yield pending.result() + except StopAsyncIteration: + return + finally: + # A client disconnect arrives as GeneratorExit at one of the + # yields above; without this the in-flight pull is orphaned + # and keeps running after the response is gone. Waiting for + # the cancellation to settle also releases the inner + # generator, which cannot be closed while a pull is in flight. + if not pending.done(): + pending.cancel() + await asyncio.wait({pending}) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() diff --git a/app/src/components/chat/ChatInterface.tsx b/app/src/components/chat/ChatInterface.tsx index 6bfd0718..c7c0cf0f 100644 --- a/app/src/components/chat/ChatInterface.tsx +++ b/app/src/components/chat/ChatInterface.tsx @@ -236,8 +236,11 @@ const ChatInterface = ({ setTimeout(() => scrollToBottom(), 50); } - // Add SQL query message with analysis info (even if SQL is empty) - if (sqlQuery !== undefined || Object.keys(analysisInfo).length > 0) { + // Add SQL query message with analysis info. sqlQuery is initialized to + // "" and is never undefined, so the old `!== undefined` guard was always + // true and painted an empty "Query Analysis" card on runs that failed + // before any SQL was generated (incident 2026-07-29). + if (sqlQuery || Object.keys(analysisInfo).length > 0) { const sqlMessage: ChatMessageData = { id: (Date.now() + 2).toString(), type: "sql-query", diff --git a/tests/test_stream_keepalive.py b/tests/test_stream_keepalive.py new file mode 100644 index 00000000..73c0468e --- /dev/null +++ b/tests/test_stream_keepalive.py @@ -0,0 +1,103 @@ +"""Tests for the streaming keepalive wrapper. + +The 2026-07-29 demo failure was a stream that emitted nothing for the whole +SQL-generation phase and was severed mid-body. ``_with_keepalive`` keeps bytes +flowing through those silent gaps. +""" + +import asyncio + +import pytest + +from api.core.pipeline import MESSAGE_DELIMITER +from api.routes.streaming import with_keepalive + + +async def _collect(agen): + return [chunk async for chunk in agen] + + +@pytest.mark.unit +async def test_passes_chunks_through_unchanged(): + """A stream that never goes idle is forwarded verbatim.""" + async def source(): + yield "a" + yield "b" + + assert await _collect(with_keepalive(source(), interval=5.0)) == ["a", "b"] + + +@pytest.mark.unit +async def test_empty_stream_yields_nothing(): + async def source(): + return + yield # pragma: no cover - never reached + + assert await _collect(with_keepalive(source(), interval=5.0)) == [] + + +@pytest.mark.unit +async def test_emits_keepalive_during_a_silent_gap(): + """A slow producer gets bare delimiters until its next real chunk.""" + async def source(): + yield "first" + await asyncio.sleep(0.25) + yield "second" + + chunks = await _collect(with_keepalive(source(), interval=0.05)) + + assert chunks[0] == "first" + assert chunks[-1] == "second" + keepalives = chunks[1:-1] + assert keepalives, "expected at least one keepalive during the gap" + assert set(keepalives) == {MESSAGE_DELIMITER} + + +@pytest.mark.unit +async def test_keepalive_is_an_empty_part_for_the_client(): + """A bare delimiter splits into empty parts, which the client skips. + + This is what makes the keepalive backward compatible: no new message type + and no client change. Mirrors the parser in app/src/services/chat.ts. + """ + payload = "".join([MESSAGE_DELIMITER, MESSAGE_DELIMITER]) + parts = [p for p in payload.split(MESSAGE_DELIMITER) if p.strip()] + assert parts == [] + + +@pytest.mark.unit +async def test_propagates_producer_exception(): + """Pipeline failures must still reach the route's error handler.""" + async def source(): + yield "a" + raise RuntimeError("pipeline exploded") + + with pytest.raises(RuntimeError, match="pipeline exploded"): + await _collect(with_keepalive(source(), interval=5.0)) + + +@pytest.mark.unit +async def test_close_mid_gap_tears_down_the_inner_stream(): + """A client disconnect during a silent gap must not orphan the pull. + + The inner generator's ``finally`` running is the observable proof that the + wrapper closed it rather than leaving it pending on the loop. + """ + closed = asyncio.Event() + + async def source(): + try: + yield "first" + await asyncio.sleep(60) # the silent gap; never completes + yield "unreachable" # pragma: no cover + finally: + closed.set() + + agen = with_keepalive(source(), interval=0.05) + assert await agen.__anext__() == "first" + # The next pull enters the gap, so this returns a keepalive, not a chunk. + assert await agen.__anext__() == MESSAGE_DELIMITER + + await agen.aclose() + + await asyncio.wait_for(closed.wait(), timeout=1) From 1c352084757ff34b03a6e5cac2c9efb2813acba9 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Tue, 18 Aug 2026 16:33:05 +0300 Subject: [PATCH 02/25] fix(agents): pin the LLM retry budget so the timeout is a real ceiling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Simulating a hung provider (a local server that accepts the request and never replies) showed the timeout aborts, but far later than configured: a 3s LLM_TIMEOUT took 10.81s to fail, because `timeout` is per attempt and both the provider SDK and litellm apply their own retry loops on top. Extrapolated to the 90s default, worst case was ~270s — long enough to defeat the point of having a timeout. Pin the budget: `max_retries` comes from the new LLM_MAX_RETRIES (default 1) and litellm's outer `num_retries` loop is disabled, so the two do not multiply. Measured after the change: the same hung provider fails in 3.19s against a 3s timeout. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- .env.example | 5 +++++ api/agents/utils.py | 9 +++++++-- api/config.py | 9 +++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/.env.example b/.env.example index 6ef6ec98..07169f99 100644 --- a/.env.example +++ b/.env.example @@ -82,6 +82,11 @@ FALKORDB_URL=redis://localhost:6379/0 # REQUIRED - change to your FalkorDB URL # # Calls slower than this are logged at WARNING (default 20). # LLM_SLOW_CALL_THRESHOLD=20 +# +# Retry budget per LLM call (default 1). LLM_TIMEOUT applies per attempt, so +# this is pinned rather than left to the provider SDK and litellm defaults, +# which each retry and together multiply the effective ceiling. +# LLM_MAX_RETRIES=1 # OpenAI - uses openai/gpt-4.1 and openai/text-embedding-ada-002 # OPENAI_API_KEY=your_openai_api_key diff --git a/api/agents/utils.py b/api/agents/utils.py index 68954eae..595d123d 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -14,8 +14,8 @@ def run_completion(messages: List[Dict[str, str]], custom_model: str = None, **kwargs) -> str: """Run an LLM completion with optional custom model/key overrides. - Applies ``Config.LLM_TIMEOUT`` unless the caller passes an explicit - ``timeout``, and logs the call duration. Both exist because the 2026-07-29 + Applies ``Config.LLM_TIMEOUT`` per attempt and a pinned retry budget + unless the caller overrides them, and logs the call duration. Both exist because the 2026-07-29 demo failure was an LLM call that stalled with no timeout and left no trace of how long it ran. ``label`` names the caller in those log lines and is not forwarded to the provider. @@ -27,6 +27,11 @@ def run_completion(messages: List[Dict[str, str]], custom_model: str = None, "messages": messages, "top_p": 1, "timeout": Config.LLM_TIMEOUT, + # ``timeout`` is per attempt, so the retry budget has to be pinned too + # or the effective ceiling becomes a multiple of it. litellm's outer + # retry loop is disabled in favour of the SDK-level count. + "max_retries": Config.LLM_MAX_RETRIES, + "num_retries": 0, **kwargs, } diff --git a/api/config.py b/api/config.py index a29c17f0..cb66031c 100644 --- a/api/config.py +++ b/api/config.py @@ -145,6 +145,15 @@ class Config: # pylint: disable-next=invalid-name LLM_SLOW_CALL_THRESHOLD: float = float(os.getenv("LLM_SLOW_CALL_THRESHOLD", "20")) + # Retry budget for a single agent LLM call. Kept explicit because the + # provider SDK and litellm each have their own retry loop, and leaving + # both at their defaults multiplies the effective ceiling (measured: a + # 3s timeout took 10.8s to fail). Applied as the SDK-level retry count + # with litellm's outer loop disabled, so the worst case stays close to + # LLM_TIMEOUT rather than a multiple of it. + # pylint: disable-next=invalid-name + LLM_MAX_RETRIES: int = int(os.getenv("LLM_MAX_RETRIES", "1")) + DB_MAX_DISTINCT: int = 100 # pylint: disable=invalid-name DB_UNIQUENESS_THRESHOLD: float = 0.5 # pylint: disable=invalid-name SHORT_MEMORY_LENGTH = 5 # Maximum number of questions to keep in short-term memory From bc50d894c76484076095d19e717d10ca6161226c Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Tue, 18 Aug 2026 16:58:31 +0300 Subject: [PATCH 03/25] test(e2e): off-topic query should show no SQL card at all MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The off-topic test asserted that the SQL card *is* visible with no SQL behind it, which encoded the phantom "Query Analysis" card from the 2026-07-29 incident rather than guarding against it. An off-topic query never reaches SQL generation: the pipeline emits only `reasoning_step` and `followup_questions`, no `sql_query` event. Since `analysisInfo` is populated solely in the `sql_query` branch, the card had nothing to render — no SQL, and no explanation either, because `isValid` defaults to true when unset. It drew a bare header. The off-topic reason already reaches the user as a normal AI message, which the test still asserts. Verified the event sequence against the real `run_query` pipeline with the relevancy agent returning Off-topic. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- e2e/tests/chat.spec.ts | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/e2e/tests/chat.spec.ts b/e2e/tests/chat.spec.ts index 305ef99f..ef43f033 100644 --- a/e2e/tests/chat.spec.ts +++ b/e2e/tests/chat.spec.ts @@ -87,11 +87,16 @@ test.describe('Chat Feature Tests', () => { const processingComplete = await homePage.waitForProcessingToComplete(); expect(processingComplete).toBeTruthy(); - // Verify Query Analysis message appears (but without actual SQL) + // Verify NO SQL card at all. An off-topic query never reaches SQL + // generation, so there is no SQL and no analysis to show — the card would + // render as a bare "Query Analysis" header with nothing under it. The + // off-topic explanation reaches the user as a normal AI message instead + // (asserted below). This previously asserted the empty card was visible, + // which masked the phantom card seen in the 2026-07-29 incident. const sqlMessageVisible = await homePage.isSQLQueryMessageVisible(); - expect(sqlMessageVisible).toBeTruthy(); + expect(sqlMessageVisible).toBeFalsy(); - // Verify NO actual SQL content (should say "Query Analysis" or "Off topic") + // And therefore no SQL content anywhere. const hasSQLContent = await homePage.verifySQLQueryContains("SELECT"); expect(hasSQLContent).toBeFalsy(); From 779c7ec88bc812a312c4e95fa131d6d78a5154c3 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Tue, 18 Aug 2026 17:13:59 +0300 Subject: [PATCH 04/25] fix: address PR #714 review feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Six findings from the Copilot and CodeRabbit reviews: - The memory path had the same blocking-call bug this PR fixes elsewhere: `update_user_information` and `summarize_conversation` are `async` but called `litellm.completion` synchronously, and they run as detached tasks via `save_memory_background` — so they could stall unrelated streaming responses. Both now go through `run_completion` inside `asyncio.to_thread`, which also gives them the shared timeout and retry bounds. (Copilot) - The render guard still had a hole: `analysisInfo` is built with all five keys defined unconditionally, so `Object.keys(...).length > 0` was always true once any `sql_query` event arrived, even with every value undefined. Check the values instead, and trim the SQL before rendering. (CodeRabbit) - The off-topic E2E assertion used `isSQLQueryMessageVisible()`, which catches locator errors and returns false, so it would pass on a broken selector. Use a strict `toHaveCount(0)` web-first assertion via a new public `sqlQueryCard` accessor, matching the existing `confirmationDialog` precedent. (CodeRabbit) - `AZURE_API_VERSION` examples in README.md and examples/README.md still showed 2024-12-01-preview, which the Responses API rejects. (CodeRabbit) - `custom_model` / `custom_api_key` annotated `str | None`. (CodeRabbit) - Test module docstring referred to `_with_keepalive`; the exported name is `with_keepalive`. (Copilot) Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 2 +- api/agents/utils.py | 4 +-- api/memory/graphiti_tool.py | 32 +++++++++++------------ app/src/components/chat/ChatInterface.tsx | 20 +++++++++----- e2e/logic/pom/homePage.ts | 9 +++++++ e2e/tests/chat.spec.ts | 6 +++-- examples/README.md | 2 +- tests/test_stream_keepalive.py | 2 +- 8 files changed, 47 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 13123df6..fa295322 100644 --- a/README.md +++ b/README.md @@ -502,7 +502,7 @@ docker run -p 5000:5000 -it \ -e FASTAPI_SECRET_KEY=your_secret_key \ -e AZURE_API_KEY=your_azure_api_key \ -e AZURE_API_BASE=https://your-resource.openai.azure.com/ \ - -e AZURE_API_VERSION=2024-12-01-preview \ + -e AZURE_API_VERSION=2025-03-01-preview \ falkordb/queryweaver ``` diff --git a/api/agents/utils.py b/api/agents/utils.py index 595d123d..aedc0099 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -9,8 +9,8 @@ from api.config import Config -def run_completion(messages: List[Dict[str, str]], custom_model: str = None, - custom_api_key: str = None, *, label: str = "llm", +def run_completion(messages: List[Dict[str, str]], custom_model: str | None = None, + custom_api_key: str | None = None, *, label: str = "llm", **kwargs) -> str: """Run an LLM completion with optional custom model/key overrides. diff --git a/api/memory/graphiti_tool.py b/api/memory/graphiti_tool.py index 96665c6d..f40551ec 100644 --- a/api/memory/graphiti_tool.py +++ b/api/memory/graphiti_tool.py @@ -27,7 +27,7 @@ from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF -from litellm import completion +from api.agents.utils import run_completion def extract_embedding_model_name(full_model_name: str) -> str: @@ -261,14 +261,14 @@ async def update_user_information(self, conversation: Dict[str, Any], history: T messages.append({"role": "user", "content": query}) messages.append({"role": "assistant", "content": result}) messages.append({"role": "user", "content": prompt}) - response = completion( - model=Config.COMPLETION_MODEL, - messages=messages, - temperature=0.1 - ) - - # Parse the direct text response (no JSON parsing needed) - content = response.choices[0].message.content.strip() + # Synchronous LLM call inside an async method that runs as a + # detached task: calling it directly would block the event loop + # and stall unrelated streaming responses. ``run_completion`` also + # applies the shared timeout and retry bounds. + content = (await asyncio.to_thread( + run_completion, messages, label="memory.user_summary", + temperature=0.1, + )).strip() query = """ MATCH (u:Entity {name: $user_id}) SET u.summary = $summary @@ -744,14 +744,12 @@ async def summarize_conversation(self, conversation: Dict[str, Any], history: Li messages.append({"role": "user", "content": query}) messages.append({"role": "assistant", "content": result}) messages.append({"role": "user", "content": prompt}) - response = completion( - model=Config.COMPLETION_MODEL, - messages=messages, - temperature=0.1 - ) - - # Parse the direct text response (no JSON parsing needed) - content = response.choices[0].message.content.strip() + # Same reasoning as ``update_user_information`` above: off-loop, + # with the shared timeout and retry bounds. + content = (await asyncio.to_thread( + run_completion, messages, label="memory.conversation_summary", + temperature=0.1, + )).strip() return { "database_summary": content } diff --git a/app/src/components/chat/ChatInterface.tsx b/app/src/components/chat/ChatInterface.tsx index c7c0cf0f..168d1a3b 100644 --- a/app/src/components/chat/ChatInterface.tsx +++ b/app/src/components/chat/ChatInterface.tsx @@ -236,15 +236,23 @@ const ChatInterface = ({ setTimeout(() => scrollToBottom(), 50); } - // Add SQL query message with analysis info. sqlQuery is initialized to - // "" and is never undefined, so the old `!== undefined` guard was always - // true and painted an empty "Query Analysis" card on runs that failed - // before any SQL was generated (incident 2026-07-29). - if (sqlQuery || Object.keys(analysisInfo).length > 0) { + // Render the SQL card only when there is genuinely something to show. + // Two traps here, both of which produced the empty "Query Analysis" + // card seen in the 2026-07-29 incident: + // - sqlQuery is initialized to "" and never undefined, so the original + // `sqlQuery !== undefined` guard was always true. + // - analysisInfo is built with all five keys defined unconditionally, + // so counting keys is always > 0 once any sql_query event arrives, + // even when every value is undefined. + const trimmedSqlQuery = sqlQuery.trim(); + const hasAnalysisInfo = Object.values(analysisInfo).some( + value => value !== undefined && value !== null && value !== '' + ); + if (trimmedSqlQuery || hasAnalysisInfo) { const sqlMessage: ChatMessageData = { id: (Date.now() + 2).toString(), type: "sql-query", - content: sqlQuery, + content: trimmedSqlQuery, analysisInfo: analysisInfo, timestamp: new Date(), }; diff --git a/e2e/logic/pom/homePage.ts b/e2e/logic/pom/homePage.ts index 6efecc32..64dd886f 100644 --- a/e2e/logic/pom/homePage.ts +++ b/e2e/logic/pom/homePage.ts @@ -113,6 +113,15 @@ export class HomePage extends BasePage { return this.confirmationMessage; } + /** + * Public accessor for the SQL-card locator, so tests can make strict + * web-first assertions instead of relying on the boolean helpers, which + * swallow locator errors and would pass on a broken selector. + */ + get sqlQueryCard(): Locator { + return this.sqlQueryMessage; + } + private get confirmationConfirmBtn(): Locator { return this.page.getByTestId("confirmation-confirm-button"); } diff --git a/e2e/tests/chat.spec.ts b/e2e/tests/chat.spec.ts index ef43f033..0f65a214 100644 --- a/e2e/tests/chat.spec.ts +++ b/e2e/tests/chat.spec.ts @@ -93,8 +93,10 @@ test.describe('Chat Feature Tests', () => { // off-topic explanation reaches the user as a normal AI message instead // (asserted below). This previously asserted the empty card was visible, // which masked the phantom card seen in the 2026-07-29 incident. - const sqlMessageVisible = await homePage.isSQLQueryMessageVisible(); - expect(sqlMessageVisible).toBeFalsy(); + // Strict web-first assertion: toHaveCount(0) waits for the final DOM + // state and fails on a broken selector, unlike the boolean helper which + // catches locator errors and returns false. + await expect(homePage.sqlQueryCard).toHaveCount(0); // And therefore no SQL content anywhere. const hasSQLContent = await homePage.verifySQLQueryContains("SELECT"); diff --git a/examples/README.md b/examples/README.md index 4391cff4..a4347c7e 100644 --- a/examples/README.md +++ b/examples/README.md @@ -86,7 +86,7 @@ For Azure OpenAI: ```bash export AZURE_API_KEY=... export AZURE_API_BASE=https://.openai.azure.com/ -export AZURE_API_VERSION=2024-12-01-preview +export AZURE_API_VERSION=2025-03-01-preview ``` Other supported providers: `GEMINI_API_KEY`, `ANTHROPIC_API_KEY`, diff --git a/tests/test_stream_keepalive.py b/tests/test_stream_keepalive.py index 73c0468e..84c0b23b 100644 --- a/tests/test_stream_keepalive.py +++ b/tests/test_stream_keepalive.py @@ -1,7 +1,7 @@ """Tests for the streaming keepalive wrapper. The 2026-07-29 demo failure was a stream that emitted nothing for the whole -SQL-generation phase and was severed mid-body. ``_with_keepalive`` keeps bytes +SQL-generation phase and was severed mid-body. ``with_keepalive`` keeps bytes flowing through those silent gaps. """ From 4d633cf81443bd43c907b0cc01a041ef4883ef26 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Wed, 19 Aug 2026 10:21:04 +0300 Subject: [PATCH 05/25] fix: offload the last two blocking calls in the query path (PR #714 review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both findings from @Naseem77 are valid, and they matter more than "two more instances of the same pattern": a keepalive cannot be written while the event loop is blocked, so these two calls could defeat the keepalive this PR adds. - `api/graph.py` `find()` called litellm and the embedding provider synchronously before its first await, while being launched via `asyncio.create_task`. That made its concurrency with the relevancy agent illusory and blocked the loop — and it is the call that logs "Calling LLM to find relevant tables/columns", the last line before the stall in the 2026-07-29 logs. Now offloaded via `asyncio.to_thread` and routed through `run_completion`, so it also picks up the shared timeout and duration logging. The embedding call is offloaded too. - `loader_class.execute_sql_query` ran on the loop in both `run_query` and `run_confirmed`. A slow query blocked every other request and stopped keepalives on its own stream. Both now offloaded. The third call site, inside `_run_sql`, already runs within the healer's thread and is left synchronous. Verified with the incident harness. With the keepalive enabled but these calls back on the loop, a 12s stall still severs the stream and delivers **zero** keepalives. With them offloaded, keepalives flow every 2s through the whole execution phase and the query completes. Starvation probe during a slow query: 63 requests served, 0.00s worst latency. All graph queries on this path were already using the async client and needed no change. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/core/text2sql.py | 12 ++++++++++-- api/graph.py | 22 +++++++++++++++------- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/api/core/text2sql.py b/api/core/text2sql.py index 5f9fb73d..ba0d4c8c 100644 --- a/api/core/text2sql.py +++ b/api/core/text2sql.py @@ -519,7 +519,12 @@ async def run_query( # pylint: disable=too-many-locals,too-many-branches,too-ma try: try: - query_results = loader_class.execute_sql_query(sql_query, db_url) + # Off-loop: driver execution is synchronous, and a slow query would + # otherwise block every other request and stop keepalives from + # flushing on this one. + query_results = await asyncio.to_thread( + loader_class.execute_sql_query, sql_query, db_url, + ) except Exception as exec_error: # pylint: disable=broad-exception-caught yield { "type": "reasoning_step", @@ -753,7 +758,10 @@ async def run_confirmed( # pylint: disable=too-many-locals,too-many-branches,to is_schema_modifying, operation_type = check_schema_modification( sql_query, loader_class, ) - query_results = loader_class.execute_sql_query(sql_query, db_url) + # Off-loop, as in ``run_query`` above. + query_results = await asyncio.to_thread( + loader_class.execute_sql_query, sql_query, db_url, + ) yield {"type": "query_result", "data": query_results} if is_schema_modifying: diff --git a/api/graph.py b/api/graph.py index 27e8ce05..cc1858b8 100644 --- a/api/graph.py +++ b/api/graph.py @@ -6,9 +6,9 @@ from itertools import combinations from typing import Any, Dict, List -from litellm import completion from pydantic import BaseModel +from api.agents.utils import run_completion from api.config import Config from api.core.db_resolver import resolve_db @@ -300,10 +300,14 @@ async def find( # pylint: disable=too-many-locals logging.info("Calling LLM to find relevant tables/columns for query") - completion_result = completion( - model=Config.COMPLETION_MODEL, - response_format=Descriptions, - messages=[ + # Both this LLM call and the embedding call below are synchronous network + # calls, and this coroutine is launched with ``asyncio.create_task``. Run + # directly they would block the event loop before the first await, which + # makes that "concurrency" illusory and prevents any stream from flushing + # keepalives — the exact point where the 2026-07-29 demo queries stalled. + completion_content = await asyncio.to_thread( + run_completion, + [ { "role": "system", "content": Config.FIND_SYSTEM_PROMPT.format( @@ -318,17 +322,21 @@ async def find( # pylint: disable=too-many-locals }) }, ], + label="find.tables", + response_format=Descriptions, temperature=0, ) - json_data = json.loads(completion_result.choices[0].message.content) + json_data = json.loads(completion_content) descriptions = Descriptions(**json_data) descriptions_text = ([desc.description for desc in descriptions.tables_descriptions] + [desc.description for desc in descriptions.columns_descriptions]) if not descriptions_text: return [] - embedding_results = Config.EMBEDDING_MODEL.embed(descriptions_text) + embedding_results = await asyncio.to_thread( + Config.EMBEDDING_MODEL.embed, descriptions_text, + ) # Split embeddings back into table and column embeddings table_embeddings = embedding_results[:len(descriptions.tables_descriptions)] From 81e3b40539a0f1e2d4730c9ea6ec37adcff1c7f0 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Wed, 19 Aug 2026 13:46:50 +0300 Subject: [PATCH 06/25] fix(streaming): remove awaiting teardown; add DB timeouts and 3-stage idle tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the remaining items from @Naseem77's review. Items 1 and 2 of his list (offload table-finding and SQL execution) landed in 4d633cf, ten minutes after that review was written against 779c7ec. **Keepalive teardown.** Rewrote `with_keepalive` so the producer runs as a task feeding a queue, instead of this generator racing `anext` against a timeout. The previous version cleaned up by awaiting a cancellation and then calling `aclose()` on the inner generator; once cancellation is pending an `await` re-raises immediately, which could leave that `aclose()` racing an in-flight pull — the `asynchronous generator is already running` signature. Cleanup is now a single non-awaiting `cancel()`, and the inner stream is consumed by a plain `async for` so its closure follows ordinary task cancellation. Note: I could not reproduce that error locally — abrupt ASGI disconnect, task cancellation mid-gap, a 60-step sweep of cancellation timings, and teardown during a non-cancellable `to_thread` call all completed cleanly on both the old and new code. The rewrite removes the construct that produces that signature rather than being verified against a reproduction. **DB timeouts**, bounding execution now that it runs in a worker thread that cannot be cancelled: `DB_CONNECT_TIMEOUT` (10s) and `DB_STATEMENT_TIMEOUT` (60s), applied in `execute_sql_query` for PostgreSQL (`connect_timeout` plus a server-side `statement_timeout`), MySQL (connect/read/write timeouts) and Snowflake (login/network timeouts plus `STATEMENT_TIMEOUT_IN_SECONDS`). Scoped to query execution, leaving the schema-load path unchanged. Loader values use `setdefault` so a URL-supplied value still wins. **Tests.** `tests/test_stream_idle_timeout.py` drives the real `run_query` through the real serializer and asserts the stream never idles longer than the keepalive interval, with the stall injected into the analysis, table-finding and SQL-execution stages in turn. `tests/test_find_offloading.py` asserts `api.graph.find` keeps the loop responsive. Both were checked against injected regressions: putting the analysis and SQL calls back on the loop fails with "no keepalive during the ... stall", and un-offloading `find` fails with "event loop was starved: 1 ticks in 1.20s (expected roughly 60)". Two more keepalive teardown tests cover cancellation timing and producer cleanup. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- .env.example | 7 ++ api/config.py | 8 ++ api/loaders/mysql_loader.py | 8 ++ api/loaders/postgres_loader.py | 12 +- api/loaders/snowflake_loader.py | 10 ++ api/routes/streaming.py | 56 +++++---- tests/test_find_offloading.py | 98 +++++++++++++++ tests/test_stream_idle_timeout.py | 196 ++++++++++++++++++++++++++++++ tests/test_stream_keepalive.py | 67 ++++++++++ 9 files changed, 439 insertions(+), 23 deletions(-) create mode 100644 tests/test_find_offloading.py create mode 100644 tests/test_stream_idle_timeout.py diff --git a/.env.example b/.env.example index 07169f99..0cd1bfd6 100644 --- a/.env.example +++ b/.env.example @@ -87,6 +87,13 @@ FALKORDB_URL=redis://localhost:6379/0 # REQUIRED - change to your FalkorDB URL # this is pinned rather than left to the provider SDK and litellm defaults, # which each retry and together multiply the effective ceiling. # LLM_MAX_RETRIES=1 +# +# Bounds for executing a user query against the target database. Offloading +# execution to a thread keeps the event loop free, but only these bound how +# long the query itself may run (a thread blocked in a socket read cannot be +# cancelled from Python). Seconds. +# DB_CONNECT_TIMEOUT=10 +# DB_STATEMENT_TIMEOUT=60 # OpenAI - uses openai/gpt-4.1 and openai/text-embedding-ada-002 # OPENAI_API_KEY=your_openai_api_key diff --git a/api/config.py b/api/config.py index cb66031c..95b04cbb 100644 --- a/api/config.py +++ b/api/config.py @@ -154,6 +154,14 @@ class Config: # pylint: disable-next=invalid-name LLM_MAX_RETRIES: int = int(os.getenv("LLM_MAX_RETRIES", "1")) + # Bounds for user-query execution against the target database. Offloading + # execution to a thread stops a slow query from blocking other requests, + # but nothing bounds how long the query itself runs without these. + # pylint: disable-next=invalid-name + DB_CONNECT_TIMEOUT: int = int(os.getenv("DB_CONNECT_TIMEOUT", "10")) + # pylint: disable-next=invalid-name + DB_STATEMENT_TIMEOUT: int = int(os.getenv("DB_STATEMENT_TIMEOUT", "60")) + DB_MAX_DISTINCT: int = 100 # pylint: disable=invalid-name DB_UNIQUENESS_THRESHOLD: float = 0.5 # pylint: disable=invalid-name SHORT_MEMORY_LENGTH = 5 # Maximum number of questions to keep in short-term memory diff --git a/api/loaders/mysql_loader.py b/api/loaders/mysql_loader.py index 2e8b40fa..2246df29 100644 --- a/api/loaders/mysql_loader.py +++ b/api/loaders/mysql_loader.py @@ -11,6 +11,7 @@ from pymysql.cursors import DictCursor +from api.config import Config from api.loaders.base_loader import BaseLoader from api.loaders.graph_loader import load_to_graph @@ -512,6 +513,13 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: # Parse connection URL conn_params = MySQLLoader._parse_mysql_url(db_url) + # Bound connect and socket waits so a hung server cannot pin this + # worker thread indefinitely. setdefault so a URL-supplied value + # still wins. + conn_params.setdefault("connect_timeout", Config.DB_CONNECT_TIMEOUT) + conn_params.setdefault("read_timeout", Config.DB_STATEMENT_TIMEOUT) + conn_params.setdefault("write_timeout", Config.DB_STATEMENT_TIMEOUT) + # Connect to MySQL database conn = pymysql.connect(**conn_params) cursor = conn.cursor(DictCursor) diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index c5dff6fa..711322a5 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -11,6 +11,7 @@ from psycopg2 import sql import tqdm +from api.config import Config from api.loaders.base_loader import BaseLoader # pylint: disable=import-error from api.loaders.graph_loader import load_to_graph # pylint: disable=import-error @@ -550,8 +551,15 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: List of dictionaries containing the query results """ try: - # Connect to PostgreSQL database - conn = psycopg2.connect(db_url) + # Bound both connection and execution. Offloading this call to a + # thread keeps the event loop free, but only a server-side + # statement_timeout bounds the query itself — and a thread blocked + # in a socket read cannot be cancelled from Python. + conn = psycopg2.connect( + db_url, + connect_timeout=Config.DB_CONNECT_TIMEOUT, + options=f"-c statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}", + ) cursor = conn.cursor() # Execute the SQL query diff --git a/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index 7685daa9..35bbecba 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -15,6 +15,7 @@ import snowflake.connector from snowflake.connector import DictCursor +from api.config import Config from api.loaders.base_loader import BaseLoader from api.loaders.graph_loader import load_to_graph @@ -646,6 +647,15 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: # Parse connection URL conn_params = SnowflakeLoader._parse_snowflake_url(db_url) + # Bound login, network waits and server-side statement runtime. + # setdefault so a URL-supplied value still wins. + conn_params.setdefault("login_timeout", Config.DB_CONNECT_TIMEOUT) + conn_params.setdefault("network_timeout", Config.DB_STATEMENT_TIMEOUT) + conn_params.setdefault( + "session_parameters", + {"STATEMENT_TIMEOUT_IN_SECONDS": Config.DB_STATEMENT_TIMEOUT}, + ) + # Connect to Snowflake database conn = snowflake.connector.connect(**conn_params) cursor = conn.cursor(DictCursor) diff --git a/api/routes/streaming.py b/api/routes/streaming.py index 50c6be2c..e99ea785 100644 --- a/api/routes/streaming.py +++ b/api/routes/streaming.py @@ -33,30 +33,44 @@ async def with_keepalive(chunks, interval: float = STREAM_KEEPALIVE_INTERVAL): including silent gaps introduced later. A bare delimiter splits into an empty part on the client, which every consumer's parser already skips, so this needs no protocol change and no client change. + + The producer runs as a task feeding a queue rather than having this + generator race ``anext`` against a timeout directly. That matters for + teardown: a client disconnect cancels the ASGI task, and cleanup that has + to ``await`` cannot complete once cancellation is pending. Here the only + cleanup is ``cancel()``, which never awaits, and ``chunks`` is consumed by + a plain ``async for`` so its closure follows ordinary task cancellation + instead of an ``aclose()`` racing an in-flight pull. """ - iterator = aiter(chunks) + queue: asyncio.Queue = asyncio.Queue() + finished = object() + + async def _pump(): + try: + async for chunk in chunks: + queue.put_nowait(chunk) + except asyncio.CancelledError: + raise + except BaseException as exc: # pylint: disable=broad-exception-caught + # Hand the failure to the consumer so the route's error handling + # still sees it, rather than losing it inside this task. + queue.put_nowait(exc) + else: + queue.put_nowait(finished) + + pump = asyncio.ensure_future(_pump()) try: while True: - pending = asyncio.ensure_future(anext(iterator)) try: - while True: - done, _ = await asyncio.wait({pending}, timeout=interval) - if done: - break - yield MESSAGE_DELIMITER - yield pending.result() - except StopAsyncIteration: + item = await asyncio.wait_for(queue.get(), timeout=interval) + except asyncio.TimeoutError: + yield MESSAGE_DELIMITER + continue + if item is finished: return - finally: - # A client disconnect arrives as GeneratorExit at one of the - # yields above; without this the in-flight pull is orphaned - # and keeps running after the response is gone. Waiting for - # the cancellation to settle also releases the inner - # generator, which cannot be closed while a pull is in flight. - if not pending.done(): - pending.cancel() - await asyncio.wait({pending}) + if isinstance(item, BaseException): + raise item + yield item finally: - aclose = getattr(iterator, "aclose", None) - if aclose is not None: - await aclose() + # Non-awaiting cleanup: safe even when cancellation is already pending. + pump.cancel() diff --git a/tests/test_find_offloading.py b/tests/test_find_offloading.py new file mode 100644 index 00000000..f240b55e --- /dev/null +++ b/tests/test_find_offloading.py @@ -0,0 +1,98 @@ +"""``api.graph.find`` must not block the event loop. + +``find`` performs a completion and an embedding call, both synchronous network +calls, and it is launched with ``asyncio.create_task`` alongside the relevancy +agent. Running them on the loop made that concurrency illusory and — more +importantly — stopped any stream from writing keepalives, since a blocked loop +cannot flush bytes. This is the call that logs "Calling LLM to find relevant +tables/columns", the last line before the stall in the 2026-07-29 logs. +""" + +import asyncio +import json +import time +import types + +import pytest + +import api.graph as graph_module + +STALL = 0.6 +TICK = 0.02 + + +class _FakeResult: + result_set: list = [] + + +class _FakeGraph: + async def query(self, query, params=None, timeout=None): + return _FakeResult() + + +class _FakeDB: + def select_graph(self, graph_id): + return _FakeGraph() + + +@pytest.fixture(name="slow_find_deps") +def _slow_find_deps(monkeypatch): + """Make find()'s two network calls slow and synchronous.""" + descriptions = json.dumps({ + "tables_descriptions": [ + {"name": "accounts", "description": "customer accounts"} + ], + "columns_descriptions": [ + {"name": "name", "description": "account name"} + ], + }) + + def slow_completion(*args, **kwargs): + time.sleep(STALL) # blocking, like the real provider call + return descriptions + + def slow_embed(texts): + time.sleep(STALL) # blocking, like the real embedding call + return [[0.0, 0.1, 0.2] for _ in texts] + + monkeypatch.setattr(graph_module, "run_completion", slow_completion) + monkeypatch.setattr(graph_module, "resolve_db", lambda db: _FakeDB()) + monkeypatch.setattr( + graph_module.Config, "EMBEDDING_MODEL", + types.SimpleNamespace(embed=slow_embed), raising=False, + ) + + +@pytest.mark.unit +async def test_find_does_not_block_the_event_loop(slow_find_deps): + """A ticker must keep running while find() is in its blocking calls.""" + ticks = [] + stop = asyncio.Event() + + async def ticker(): + while not stop.is_set(): + ticks.append(time.monotonic()) + await asyncio.sleep(TICK) + + ticker_task = asyncio.ensure_future(ticker()) + started = time.monotonic() + await graph_module.find("g", ["show me customers"], "CRM demo.") + elapsed = time.monotonic() - started + stop.set() + await ticker_task + + # Both blocking calls ran, so this took at least 2 * STALL. + assert elapsed >= STALL * 2 * 0.9, f"stalls did not run (elapsed {elapsed:.2f}s)" + + # The loop stayed responsive throughout: with both calls offloaded the + # ticker keeps firing. If they ran on the loop it would be starved and + # produce only a couple of ticks. + expected = elapsed / TICK + assert len(ticks) > expected * 0.4, ( + f"event loop was starved: {len(ticks)} ticks in {elapsed:.2f}s " + f"(expected roughly {expected:.0f})" + ) + + # And the longest gap between ticks stays far below the stall duration. + gaps = [b - a for a, b in zip(ticks, ticks[1:])] + assert max(gaps) < STALL / 2, f"loop blocked for {max(gaps):.2f}s" diff --git a/tests/test_stream_idle_timeout.py b/tests/test_stream_idle_timeout.py new file mode 100644 index 00000000..e3f5aad3 --- /dev/null +++ b/tests/test_stream_idle_timeout.py @@ -0,0 +1,196 @@ +"""End-to-end idle-timeout coverage for the three slow stages. + +The 2026-07-29 demo failure was a stream that went silent long enough for an +intermediary to sever it. Silence can come from any stage that performs a slow +blocking call, and a keepalive cannot be written while the event loop is +blocked — so covering one stage is not enough. These tests drive the real +``run_query`` generator through the real route-layer serialization and assert +that the stream never goes idle longer than the keepalive interval, with the +stall injected into each stage in turn: + + * analysis -> AnalysisAgent.get_analysis + * table finding -> api.graph.find + * SQL execution -> loader.execute_sql_query + +Each stage is stalled with a genuinely blocking ``time.sleep`` so a regression +that puts the call back on the event loop shows up as a missing keepalive +rather than passing silently. +""" + +import asyncio +import json +import time + +import pytest + +from api.core.pipeline import MESSAGE_DELIMITER +from api.routes.streaming import with_keepalive + +STALL = 0.9 +INTERVAL = 0.15 +# Generous: the assertion is "keepalives kept flowing", not a latency budget. +MAX_IDLE = INTERVAL * 4 + +ANALYSIS = { + "sql_query": "SELECT name FROM accounts LIMIT 5", + "confidence": 0.9, + "missing_information": "", + "ambiguities": "", + "explanation": "Lists customers.", + "is_sql_translatable": True, +} +TABLES = [[ + "accounts", "Customer accounts.", {}, + [{"columnName": "name", "dataType": "text", "description": "Account name"}], +]] + + +class _Chat: + """Minimal stand-in for ChatRequest.""" + + def __init__(self, query="Show me five customers"): + self.chat = [query] + self.result = None + self.instructions = None + self.use_user_rules = False + self.use_memory = False + self.custom_api_key = None + self.custom_model = None + + +class _Loader: + stall = False + + @staticmethod + def execute_sql_query(sql, db_url): + if _Loader.stall: + time.sleep(STALL) + return [{"name": "Stark Industries"}] + + +@pytest.fixture(name="pipeline_stubs") +def _pipeline_stubs(monkeypatch): + """Stub the external seams, leaving the pipeline's own structure real.""" + import api.core.text2sql as t2s + + _Loader.stall = False + + async def fake_db_description(namespaced, db=None): + return ("CRM demo.", "postgresql://u:p@localhost:5432/demo") + + async def fake_find(namespaced, queries_history, db_description, db=None): + return TABLES + + monkeypatch.setattr(t2s, "get_db_description", fake_db_description) + monkeypatch.setattr(t2s, "find", fake_find) + monkeypatch.setattr(t2s, "get_user_rules", lambda *a, **k: None) + monkeypatch.setattr( + t2s, "get_database_type_and_loader", lambda url: ("postgresql", _Loader) + ) + monkeypatch.setattr(t2s, "check_schema_modification", lambda sql, loader: (False, None)) + monkeypatch.setattr(t2s, "detect_destructive_operation", lambda sql, db_type: (None, False)) + monkeypatch.setattr(t2s, "auto_quote_sql_identifiers", lambda sql, *a, **k: (sql, False)) + monkeypatch.setattr(t2s, "is_general_graph", lambda *a, **k: False) + monkeypatch.setattr(t2s, "save_memory_background", lambda *a, **k: None) + monkeypatch.setattr(t2s, "format_ai_response", lambda **k: "Here are five customers.") + + class _Relevancy: + def __init__(self, *a, **k): + pass + + async def get_answer(self, *a, **k): + return {"status": "On-topic", "reason": "about accounts"} + + class _Analysis: + stall = False + + def __init__(self, *a, **k): + pass + + def get_analysis(self, *a, **k): + if _Analysis.stall: + time.sleep(STALL) + return dict(ANALYSIS) + + monkeypatch.setattr(t2s, "RelevancyAgent", _Relevancy) + monkeypatch.setattr(t2s, "AnalysisAgent", _Analysis) + return {"analysis": _Analysis, "loader": _Loader, "text2sql": t2s} + + +async def _collect_gaps(t2s): + """Run the pipeline through the real serializer + keepalive, timing arrivals.""" + from api.core.text2sql import _Final + + async def serialize(gen): + async for event in gen: + if isinstance(event, _Final): + break + yield json.dumps(event) + MESSAGE_DELIMITER + + gaps, payloads, keepalives = [], [], 0 + last = time.monotonic() + stream = with_keepalive( + serialize(t2s.run_query("u", "g", _Chat())), interval=INTERVAL + ) + async for chunk in stream: + now = time.monotonic() + gaps.append(now - last) + last = now + if chunk == MESSAGE_DELIMITER: + keepalives += 1 + else: + payloads.append(chunk) + return max(gaps), keepalives, payloads + + +@pytest.mark.unit +async def test_no_stall_completes_without_idle_gap(pipeline_stubs): + max_gap, _, payloads = await _collect_gaps(pipeline_stubs["text2sql"]) + assert max_gap < MAX_IDLE, f"stream idle for {max_gap:.2f}s" + assert any('"ai_response"' in p for p in payloads) + + +@pytest.mark.unit +async def test_slow_analysis_stage_keeps_stream_alive(pipeline_stubs): + """Stage 1: the analysis LLM — the stall seen in the incident.""" + pipeline_stubs["analysis"].stall = True + try: + max_gap, keepalives, payloads = await _collect_gaps(pipeline_stubs["text2sql"]) + finally: + pipeline_stubs["analysis"].stall = False + + assert keepalives >= 2, "no keepalive during the analysis stall" + assert max_gap < MAX_IDLE, f"stream idle for {max_gap:.2f}s during analysis" + assert any('"ai_response"' in p for p in payloads) + + +@pytest.mark.unit +async def test_slow_table_finding_keeps_stream_alive(pipeline_stubs, monkeypatch): + """Stage 2: table finding — where the incident logs actually stop.""" + t2s = pipeline_stubs["text2sql"] + + async def slow_find(namespaced, queries_history, db_description, db=None): + # api.graph.find offloads its blocking LLM/embedding work; mirror that. + await asyncio.to_thread(time.sleep, STALL) + return TABLES + + monkeypatch.setattr(t2s, "find", slow_find) + max_gap, keepalives, payloads = await _collect_gaps(t2s) + + assert keepalives >= 2, "no keepalive during the table-finding stall" + assert max_gap < MAX_IDLE, f"stream idle for {max_gap:.2f}s during table finding" + assert any('"ai_response"' in p for p in payloads) + + +@pytest.mark.unit +async def test_slow_sql_execution_keeps_stream_alive(pipeline_stubs): + """Stage 3: database execution.""" + pipeline_stubs["loader"].stall = True + try: + max_gap, keepalives, payloads = await _collect_gaps(pipeline_stubs["text2sql"]) + finally: + pipeline_stubs["loader"].stall = False + + assert keepalives >= 2, "no keepalive during the SQL-execution stall" + assert max_gap < MAX_IDLE, f"stream idle for {max_gap:.2f}s during SQL execution" + assert any('"query_result"' in p for p in payloads) diff --git a/tests/test_stream_keepalive.py b/tests/test_stream_keepalive.py index 84c0b23b..217a5204 100644 --- a/tests/test_stream_keepalive.py +++ b/tests/test_stream_keepalive.py @@ -101,3 +101,70 @@ async def source(): await agen.aclose() await asyncio.wait_for(closed.wait(), timeout=1) + + +@pytest.mark.unit +async def test_cancellation_at_arbitrary_moments_never_raises(): + """Teardown must be clean no matter when cancellation lands. + + A client disconnect cancels the ASGI task at an arbitrary point. Cleanup + that awaits cannot finish once cancellation is pending, which previously + let an ``aclose()`` race an in-flight pull and raise + ``asynchronous generator is already running``. Sweep the cancellation + point across the keepalive cycle to cover that window. + """ + errors = [] + + for step in range(60): + async def source(): + yield "first" + await asyncio.sleep(5) # silent gap, pull stays in flight + yield "never" # pragma: no cover + + agen = with_keepalive(source(), interval=0.01) + + async def consume(gen): + async for _ in gen: + pass + + task = asyncio.ensure_future(consume(agen)) + # Sweep across (and past) the keepalive interval in small increments. + await asyncio.sleep(0.001 + step * 0.0005) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Starlette closes the body iterator after cancelling it. + try: + await agen.aclose() + except asyncio.CancelledError: + pass + except RuntimeError as exc: # the regression we are guarding + errors.append(f"step={step}: {exc}") + + assert not errors, "teardown raised: " + "; ".join(errors) + + +@pytest.mark.unit +async def test_producer_is_cancelled_when_consumer_stops_early(): + """Abandoning the stream must not leave the pipeline running.""" + cancelled = asyncio.Event() + + async def source(): + try: + yield "first" + await asyncio.sleep(60) + yield "never" # pragma: no cover + except asyncio.CancelledError: + cancelled.set() + raise + finally: + cancelled.set() + + agen = with_keepalive(source(), interval=0.01) + assert await agen.__anext__() == "first" + await agen.aclose() + + await asyncio.wait_for(cancelled.wait(), timeout=2) From 72dc7f52aaa98474770ab80313523ad0445e9c13 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Wed, 19 Aug 2026 14:03:05 +0300 Subject: [PATCH 07/25] fix(loaders,streaming): correct the DB timeout wiring and bound the queue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up review found the timeout work from 81e3b40 was partly ineffective and partly unsafe. Three real defects, all now covered by tests: - **Snowflake timeouts never applied.** `_parse_snowflake_url` hardcodes `login_timeout=30` / `network_timeout=60`, so the `setdefault` calls were no-ops and `DB_CONNECT_TIMEOUT=10` still produced a 30s login timeout. Assign instead. The regression test fails with `assert 30 == 10` if this reverts. - **PostgreSQL could drop URL connection options.** Keyword arguments override the DSN, so a bare `options=` replaced any URL-supplied `options`, silently losing e.g. `search_path`. Merge the URL's options and append `statement_timeout`, and leave an explicit URL value alone. Extracted to `_execution_connect_kwargs` to keep `execute_sql_query` within the lint limits. - **Comments claimed URL values win; they cannot.** `_parse_mysql_url` discards the query string entirely, so no URL-supplied timeout can reach it. Assign the configured values and say so accurately. Streaming, from the same review: - Bound the producer queue (`maxsize=1`) so a slow client cannot accumulate the whole source stream in memory. - Catch `Exception` rather than `BaseException` in the pump, so `SystemExit`, `KeyboardInterrupt` and `GeneratorExit` are no longer relayed to the consumer as stream errors. The explicit `CancelledError` re-raise is gone with it — `CancelledError` is a `BaseException`, so it already propagates untouched. Also: consistent import style and explanatory comments on the intentionally empty `except` blocks, and the `find()` test now asserts on the returned value. New: tests/test_db_execution_timeouts.py (5 tests) pinning that the configured values actually reach each driver. 233 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/loaders/mysql_loader.py | 11 ++-- api/loaders/postgres_loader.py | 36 ++++++++--- api/loaders/snowflake_loader.py | 15 +++-- api/routes/streaming.py | 18 +++--- tests/test_db_execution_timeouts.py | 92 +++++++++++++++++++++++++++++ tests/test_find_offloading.py | 5 +- tests/test_stream_idle_timeout.py | 2 +- tests/test_stream_keepalive.py | 4 ++ 8 files changed, 155 insertions(+), 28 deletions(-) create mode 100644 tests/test_db_execution_timeouts.py diff --git a/api/loaders/mysql_loader.py b/api/loaders/mysql_loader.py index 2246df29..fc922ac3 100644 --- a/api/loaders/mysql_loader.py +++ b/api/loaders/mysql_loader.py @@ -514,11 +514,12 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: conn_params = MySQLLoader._parse_mysql_url(db_url) # Bound connect and socket waits so a hung server cannot pin this - # worker thread indefinitely. setdefault so a URL-supplied value - # still wins. - conn_params.setdefault("connect_timeout", Config.DB_CONNECT_TIMEOUT) - conn_params.setdefault("read_timeout", Config.DB_STATEMENT_TIMEOUT) - conn_params.setdefault("write_timeout", Config.DB_STATEMENT_TIMEOUT) + # worker thread indefinitely. ``_parse_mysql_url`` discards the + # URL's query string, so there is no URL-supplied value to preserve + # here — these are the only timeouts in play. + conn_params["connect_timeout"] = Config.DB_CONNECT_TIMEOUT + conn_params["read_timeout"] = Config.DB_STATEMENT_TIMEOUT + conn_params["write_timeout"] = Config.DB_STATEMENT_TIMEOUT # Connect to MySQL database conn = pymysql.connect(**conn_params) diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index 711322a5..6080ae6f 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -537,6 +537,34 @@ async def refresh_graph_schema(graph_id: str, db_url: str, db=None) -> Tuple[boo logging.error(error_msg) return False, error_msg + @staticmethod + def _execution_connect_kwargs(db_url: str) -> Dict[str, Any]: + """Timeout keywords for executing a user query. + + Offloading execution to a thread keeps the event loop free, but only a + server-side ``statement_timeout`` bounds the query itself — a thread + blocked in a socket read cannot be cancelled from Python. + + Keyword arguments override values in the DSN, so anything the URL + already specifies is merged rather than replaced: a bare ``options=`` + would silently drop a URL-supplied ``search_path``. + """ + url_params = parse_qs(urlparse(db_url).query) + url_options = url_params.get("options", [""])[0] + kwargs: Dict[str, Any] = {} + + if "statement_timeout" in url_options: + options = url_options + else: + timeout_ms = Config.DB_STATEMENT_TIMEOUT * 1000 + options = f"{url_options} -c statement_timeout={timeout_ms}".strip() + if options: + kwargs["options"] = options + + if "connect_timeout" not in url_params: + kwargs["connect_timeout"] = Config.DB_CONNECT_TIMEOUT + return kwargs + @staticmethod def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: """ @@ -551,14 +579,8 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: List of dictionaries containing the query results """ try: - # Bound both connection and execution. Offloading this call to a - # thread keeps the event loop free, but only a server-side - # statement_timeout bounds the query itself — and a thread blocked - # in a socket read cannot be cancelled from Python. conn = psycopg2.connect( - db_url, - connect_timeout=Config.DB_CONNECT_TIMEOUT, - options=f"-c statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}", + db_url, **PostgresLoader._execution_connect_kwargs(db_url) ) cursor = conn.cursor() diff --git a/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index 35bbecba..4dd7d246 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -648,13 +648,16 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: conn_params = SnowflakeLoader._parse_snowflake_url(db_url) # Bound login, network waits and server-side statement runtime. - # setdefault so a URL-supplied value still wins. - conn_params.setdefault("login_timeout", Config.DB_CONNECT_TIMEOUT) - conn_params.setdefault("network_timeout", Config.DB_STATEMENT_TIMEOUT) - conn_params.setdefault( - "session_parameters", - {"STATEMENT_TIMEOUT_IN_SECONDS": Config.DB_STATEMENT_TIMEOUT}, + # ``_parse_snowflake_url`` hardcodes login_timeout/network_timeout, + # so these must be assigned — setdefault would be a no-op and the + # configured values would never apply. + conn_params["login_timeout"] = Config.DB_CONNECT_TIMEOUT + conn_params["network_timeout"] = Config.DB_STATEMENT_TIMEOUT + session_parameters = dict(conn_params.get("session_parameters") or {}) + session_parameters.setdefault( + "STATEMENT_TIMEOUT_IN_SECONDS", Config.DB_STATEMENT_TIMEOUT ) + conn_params["session_parameters"] = session_parameters # Connect to Snowflake database conn = snowflake.connector.connect(**conn_params) diff --git a/api/routes/streaming.py b/api/routes/streaming.py index e99ea785..dfdfe2af 100644 --- a/api/routes/streaming.py +++ b/api/routes/streaming.py @@ -42,21 +42,23 @@ async def with_keepalive(chunks, interval: float = STREAM_KEEPALIVE_INTERVAL): a plain ``async for`` so its closure follows ordinary task cancellation instead of an ``aclose()`` racing an in-flight pull. """ - queue: asyncio.Queue = asyncio.Queue() + # maxsize=1 keeps backpressure: without it a slow client would let the + # whole source stream accumulate in memory. + queue: asyncio.Queue = asyncio.Queue(maxsize=1) finished = object() async def _pump(): try: async for chunk in chunks: - queue.put_nowait(chunk) - except asyncio.CancelledError: - raise - except BaseException as exc: # pylint: disable=broad-exception-caught + await queue.put(chunk) + except Exception as exc: # pylint: disable=broad-exception-caught # Hand the failure to the consumer so the route's error handling # still sees it, rather than losing it inside this task. - queue.put_nowait(exc) + # CancelledError is a BaseException, so it is not caught here and + # propagates as cancellation should. + await queue.put(exc) else: - queue.put_nowait(finished) + await queue.put(finished) pump = asyncio.ensure_future(_pump()) try: @@ -68,7 +70,7 @@ async def _pump(): continue if item is finished: return - if isinstance(item, BaseException): + if isinstance(item, Exception): raise item yield item finally: diff --git a/tests/test_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py new file mode 100644 index 00000000..3a7500a3 --- /dev/null +++ b/tests/test_db_execution_timeouts.py @@ -0,0 +1,92 @@ +"""Timeout bounds applied when executing a user query. + +Query execution runs in a worker thread so it cannot block the event loop, but +a thread blocked in a socket read cannot be cancelled from Python — so the only +thing bounding a slow query is a driver/server-side timeout. These tests pin +that the configured values actually reach the driver, which is easy to get +wrong: two of the three URL parsers discard or hardcode these keys, so a +``setdefault`` silently does nothing. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +import api.core # noqa: F401 (import first: the loaders import via api.core) +from api.config import Config +from api.loaders.mysql_loader import MySQLLoader +from api.loaders.postgres_loader import PostgresLoader +from api.loaders.snowflake_loader import SnowflakeLoader + +PG_URL = "postgresql://u:p@h:5432/db" + + +@pytest.mark.unit +def test_postgres_applies_statement_and_connect_timeouts(): + kwargs = PostgresLoader._execution_connect_kwargs(PG_URL) + assert kwargs["connect_timeout"] == Config.DB_CONNECT_TIMEOUT + assert f"statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}" in kwargs["options"] + + +@pytest.mark.unit +def test_postgres_preserves_url_options(): + """A bare options= kwarg would drop a URL-supplied search_path.""" + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options=-c%20search_path%3Dfoo" + ) + assert "search_path=foo" in kwargs["options"] + assert "statement_timeout" in kwargs["options"] + + +@pytest.mark.unit +def test_postgres_url_timeouts_win(): + """An explicit value in the URL is not overridden.""" + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options=-c%20statement_timeout%3D1234" + ) + assert kwargs["options"] == "-c statement_timeout=1234" + + kwargs = PostgresLoader._execution_connect_kwargs(f"{PG_URL}?connect_timeout=3") + assert "connect_timeout" not in kwargs + + +@pytest.mark.unit +@patch("api.loaders.mysql_loader.pymysql.connect") +def test_mysql_applies_timeouts(mock_connect): + cursor = MagicMock() + cursor.description = None + cursor.rowcount = 0 + mock_connect.return_value.cursor.return_value = cursor + + MySQLLoader.execute_sql_query("SELECT 1", "mysql://u:p@h:3306/db") + + params = mock_connect.call_args.kwargs + assert params["connect_timeout"] == Config.DB_CONNECT_TIMEOUT + assert params["read_timeout"] == Config.DB_STATEMENT_TIMEOUT + assert params["write_timeout"] == Config.DB_STATEMENT_TIMEOUT + + +@pytest.mark.unit +@patch("api.loaders.snowflake_loader.snowflake.connector.connect") +def test_snowflake_overrides_parser_timeout_defaults(mock_connect): + """The parser hardcodes login_timeout=30/network_timeout=60. + + A ``setdefault`` here would be a no-op, leaving the configured values + unused — which is the bug this pins. + """ + cursor = MagicMock() + cursor.description = None + cursor.rowcount = 0 + mock_connect.return_value.cursor.return_value = cursor + + SnowflakeLoader.execute_sql_query( + "SELECT 1", "snowflake://u:p@acct/db/schema?warehouse=WH" + ) + + params = mock_connect.call_args.kwargs + assert params["login_timeout"] == Config.DB_CONNECT_TIMEOUT + assert params["network_timeout"] == Config.DB_STATEMENT_TIMEOUT + assert ( + params["session_parameters"]["STATEMENT_TIMEOUT_IN_SECONDS"] + == Config.DB_STATEMENT_TIMEOUT + ) diff --git a/tests/test_find_offloading.py b/tests/test_find_offloading.py index f240b55e..c6d369dd 100644 --- a/tests/test_find_offloading.py +++ b/tests/test_find_offloading.py @@ -76,7 +76,7 @@ async def ticker(): ticker_task = asyncio.ensure_future(ticker()) started = time.monotonic() - await graph_module.find("g", ["show me customers"], "CRM demo.") + tables = await graph_module.find("g", ["show me customers"], "CRM demo.") elapsed = time.monotonic() - started stop.set() await ticker_task @@ -96,3 +96,6 @@ async def ticker(): # And the longest gap between ticks stays far below the stall duration. gaps = [b - a for a, b in zip(ticks, ticks[1:])] assert max(gaps) < STALL / 2, f"loop blocked for {max(gaps):.2f}s" + + # The fake graph returns no rows, so the call still completes normally. + assert tables == [] diff --git a/tests/test_stream_idle_timeout.py b/tests/test_stream_idle_timeout.py index e3f5aad3..b5c8ee69 100644 --- a/tests/test_stream_idle_timeout.py +++ b/tests/test_stream_idle_timeout.py @@ -71,7 +71,7 @@ def execute_sql_query(sql, db_url): @pytest.fixture(name="pipeline_stubs") def _pipeline_stubs(monkeypatch): """Stub the external seams, leaving the pipeline's own structure real.""" - import api.core.text2sql as t2s + from api.core import text2sql as t2s _Loader.stall = False diff --git a/tests/test_stream_keepalive.py b/tests/test_stream_keepalive.py index 217a5204..a41c5469 100644 --- a/tests/test_stream_keepalive.py +++ b/tests/test_stream_keepalive.py @@ -134,12 +134,16 @@ async def consume(gen): try: await task except asyncio.CancelledError: + # Expected: we cancelled it. The assertion is about teardown, not + # about how the consumer ended. pass # Starlette closes the body iterator after cancelling it. try: await agen.aclose() except asyncio.CancelledError: + # Also acceptable: closing during cancellation may surface the + # cancellation itself. Only a RuntimeError is a defect. pass except RuntimeError as exc: # the regression we are guarding errors.append(f"step={step}: {exc}") From 07a8e59a2ec5684cde3ab9e95a759970170a91ae Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Wed, 19 Aug 2026 14:12:36 +0300 Subject: [PATCH 08/25] fix(loaders): match a real statement_timeout directive, not the bare word Two more findings from the automated reviews: - The PostgreSQL check for an existing `statement_timeout` was a substring test, so an unrelated option value such as `-c application_name=statement_timeout_probe` looked like an existing bound and the configured timeout was silently skipped. Match an actual `-c statement_timeout=` directive instead, with a regression test for the lookalike. - Dropped the side-effect-only `import api.core` from the new test by importing the loaders through `api.core.pipeline`, which initialises the package in the right order. Importing `api.loaders.postgres_loader` first hits a pre-existing circular import, so the import path here is load-bearing and now says so. 234 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/loaders/postgres_loader.py | 8 +++++++- tests/test_db_execution_timeouts.py | 22 +++++++++++++++++++--- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index 6080ae6f..b0325689 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -553,7 +553,13 @@ def _execution_connect_kwargs(db_url: str) -> Dict[str, Any]: url_options = url_params.get("options", [""])[0] kwargs: Dict[str, Any] = {} - if "statement_timeout" in url_options: + # Match an actual directive, not the bare word: a substring test would + # also hit something like ``-c application_name=statement_timeout_probe`` + # and silently skip our bound. + has_statement_timeout = re.search( + r"(?:^|\s)-c\s*statement_timeout\s*=", url_options + ) + if has_statement_timeout: options = url_options else: timeout_ms = Config.DB_STATEMENT_TIMEOUT * 1000 diff --git a/tests/test_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py index 3a7500a3..b6bb406e 100644 --- a/tests/test_db_execution_timeouts.py +++ b/tests/test_db_execution_timeouts.py @@ -12,10 +12,12 @@ import pytest -import api.core # noqa: F401 (import first: the loaders import via api.core) from api.config import Config -from api.loaders.mysql_loader import MySQLLoader -from api.loaders.postgres_loader import PostgresLoader +# Imported via api.core.pipeline: importing api.loaders.postgres_loader first +# hits a circular import (pipeline imports the loaders, the loaders import +# api.core). Going through pipeline initialises the package in the right order, +# which also makes the snowflake import below work. +from api.core.pipeline import MySQLLoader, PostgresLoader from api.loaders.snowflake_loader import SnowflakeLoader PG_URL = "postgresql://u:p@h:5432/db" @@ -50,6 +52,20 @@ def test_postgres_url_timeouts_win(): assert "connect_timeout" not in kwargs +@pytest.mark.unit +def test_postgres_ignores_a_statement_timeout_lookalike(): + """Only a real ``-c statement_timeout=`` directive counts. + + A substring test would treat this option value as an existing timeout and + silently skip the configured bound. + """ + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options=-c%20application_name%3Dstatement_timeout_probe" + ) + assert "application_name=statement_timeout_probe" in kwargs["options"] + assert f"-c statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}" in kwargs["options"] + + @pytest.mark.unit @patch("api.loaders.mysql_loader.pymysql.connect") def test_mysql_applies_timeouts(mock_connect): From 0d0257422d24f08ceca89258d0b853b8adced5ac Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Thu, 20 Aug 2026 09:50:38 +0300 Subject: [PATCH 09/25] fix: offload embeddings and schema loading; clamp URL timeout overrides MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses @Naseem77's third review. All three findings were valid. **1. Memory path blocked the loop.** Five embedding calls in `api/memory/graphiti_tool.py` (99, 127, 164, 358, 441) were synchronous inside `async def` methods, and `search_memories` runs inside the silent window before the SQL chunk. The request model defaults `use_memory` to False but `ChatInterface.tsx:54` sends `useMemory = true`, so this is on for browser traffic. Centralised behind `api/embeddings.py` (`embed_off_loop`, `embed_one_off_loop`, `vector_size_off_loop`), and `EmbeddingsModel.embed` / `get_vector_size` now carry the same timeout and pinned retry budget as the completion path, plus duration logging. **2. Schema streams blocked.** `load_to_graph` embedded per table and per column batch inline — the largest blocking stretch, and it backs the connect and refresh streams. Now offloaded, along with connect and introspection in all three loaders' `load()`. Only the connect is time-bounded there: introspecting a very large schema can legitimately outlast `DB_STATEMENT_TIMEOUT`, which is sized for user queries. **3. URLs could disable execution bounds.** `statement_timeout=0` / `connect_timeout=0` in the URL are "no limit" and overrode the configured safeguards, letting one query hold a shared worker thread indefinitely. Configured values are now maximums: a URL may tighten them, but a looser or disabled value is replaced rather than appended alongside. Tests (three new files, 12 tests): - `test_embeddings_offloading.py` — the helpers keep the loop responsive, calls are time-bounded, and a guard fails on any bare `EMBEDDING_MODEL.embed(` in the async modules. - `test_schema_load_offloading.py` — `load()` keeps the loop responsive. - `test_stream_idle_timeout.py` — a fourth stage covering memory search with `use_memory=True`. - `test_db_execution_timeouts.py` — clamping cases for disabled/looser URLs. Each was checked against an injected regression: un-offloading fails with "loop blocked for 0.32s", and reinstating an inline embed fails the call-site guard. 246 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/config.py | 39 +++++++++- api/embeddings.py | 30 ++++++++ api/loaders/graph_loader.py | 15 ++-- api/loaders/mysql_loader.py | 22 +++++- api/loaders/postgres_loader.py | 60 +++++++++++---- api/loaders/snowflake_loader.py | 18 ++++- api/memory/graphiti_tool.py | 14 ++-- tests/test_db_execution_timeouts.py | 30 +++++++- tests/test_embeddings_offloading.py | 109 +++++++++++++++++++++++++++ tests/test_schema_load_offloading.py | 98 ++++++++++++++++++++++++ tests/test_stream_idle_timeout.py | 53 +++++++++++++ 11 files changed, 451 insertions(+), 37 deletions(-) create mode 100644 api/embeddings.py create mode 100644 tests/test_embeddings_offloading.py create mode 100644 tests/test_schema_load_offloading.py diff --git a/api/config.py b/api/config.py index 95b04cbb..9f83b975 100644 --- a/api/config.py +++ b/api/config.py @@ -4,6 +4,7 @@ """ import os +import time import logging import dataclasses from typing import Union @@ -40,10 +41,25 @@ def __init__(self, model_name: str, config: dict = None): self.model_name = model_name self.config = config + def _embedding_kwargs(self) -> dict: + """Timeout and retry bounds, matching the completion path. + + These are blocking network calls, so an unbounded one pins whichever + thread runs it. ``timeout`` is per attempt, so the retry budget is + pinned too or the effective ceiling becomes a multiple of it. + """ + return { + "timeout": Config.LLM_TIMEOUT, + "max_retries": Config.LLM_MAX_RETRIES, + "num_retries": 0, + } + def embed(self, text: Union[str, list]) -> list: """ Get the embeddings of the text + Blocking: call via ``api.embeddings.embed_off_loop`` from async code. + Args: text (str|list): The text(s) to embed @@ -51,7 +67,21 @@ def embed(self, text: Union[str, list]) -> list: list: The embeddings of the text """ - embeddings = embedding(model=self.model_name, input=text) + started = time.monotonic() + try: + embeddings = embedding( + model=self.model_name, input=text, **self._embedding_kwargs() + ) + except Exception: + logging.warning( + "embed_call model=%s duration=%.2fs outcome=error", + self.model_name, time.monotonic() - started, + ) + raise + logging.info( + "embed_call model=%s duration=%.2fs outcome=ok", + self.model_name, time.monotonic() - started, + ) embeddings = [embedding["embedding"] for embedding in embeddings.data] return embeddings @@ -59,11 +89,16 @@ def get_vector_size(self) -> int: """ Get the size of the vector + Blocking: call via ``api.embeddings.vector_size_off_loop`` from async + code. + Returns: int: The size of the vector """ - response = embedding(input=["Hello World"], model=self.model_name) + response = embedding( + input=["Hello World"], model=self.model_name, **self._embedding_kwargs() + ) size = len(response.data[0]["embedding"]) return size diff --git a/api/embeddings.py b/api/embeddings.py new file mode 100644 index 00000000..82ca1aac --- /dev/null +++ b/api/embeddings.py @@ -0,0 +1,30 @@ +"""Off-loop embedding helpers. + +``EmbeddingsModel.embed`` and ``get_vector_size`` are blocking network calls. +Called directly from an async method they block the event loop, which stops +every open stream from flushing keepalives — the same failure mode as the +analysis and SQL-execution stages in the 2026-07-29 incident. Async callers go +through these helpers so the offload (and the timeout bounds inside the model) +apply in one place. +""" + +import asyncio +from typing import List, Union + +from api.config import Config + + +async def embed_off_loop(text: Union[str, list]) -> List[List[float]]: + """Embed *text* in a worker thread. Returns one vector per input.""" + return await asyncio.to_thread(Config.EMBEDDING_MODEL.embed, text) + + +async def embed_one_off_loop(text: str) -> List[float]: + """Embed a single string in a worker thread and return its vector.""" + vectors = await embed_off_loop(text) + return vectors[0] + + +async def vector_size_off_loop() -> int: + """Probe the embedding dimensionality in a worker thread.""" + return await asyncio.to_thread(Config.EMBEDDING_MODEL.get_vector_size) diff --git a/api/loaders/graph_loader.py b/api/loaders/graph_loader.py index b4f3fca5..4ff07acc 100644 --- a/api/loaders/graph_loader.py +++ b/api/loaders/graph_loader.py @@ -4,8 +4,8 @@ import tqdm -from api.config import Config from api.core.db_resolver import resolve_db +from api.embeddings import embed_off_loop, vector_size_off_loop from api.utils import generate_db_description, create_combined_description @@ -30,8 +30,11 @@ async def load_to_graph( # pylint: disable=too-many-arguments,too-many-position - db: Optional FalkorDB handle; falls back to the server singleton. """ graph = resolve_db(db).select_graph(graph_id) - embedding_model = Config.EMBEDDING_MODEL - vec_len = embedding_model.get_vector_size() + # Off-loop: these are blocking network calls, and this coroutine runs + # inside the connect/refresh streaming responses. Running them directly + # blocks the event loop, so those streams cannot emit keepalives while a + # large schema is loading. + vec_len = await vector_size_off_loop() create_combined_description(entities) @@ -70,7 +73,7 @@ async def load_to_graph( # pylint: disable=too-many-arguments,too-many-position for table_name, table_info in tqdm.tqdm(entities.items(), desc="Creating Graph Table Nodes"): table_desc = table_info["description"] - embedding_result = embedding_model.embed(table_desc) + embedding_result = await embed_off_loop(table_desc) fk = json.dumps(table_info.get("foreign_keys", [])) # Create table node @@ -109,7 +112,7 @@ async def load_to_graph( # pylint: disable=too-many-arguments,too-many-position desc=f"Creating embeddings for {table_name} columns", ): - embedding_result = embedding_model.embed(batch) + embedding_result = await embed_off_loop(batch) embed_columns.extend(embedding_result) except Exception as e: # pylint: disable=broad-exception-caught print(f"Error creating embeddings: {str(e)}") @@ -123,7 +126,7 @@ async def load_to_graph( # pylint: disable=too-many-arguments,too-many-position ): if not batch_flag: embed_columns = [] - embedding_result = embedding_model.embed(col_info["description"]) + embedding_result = await embed_off_loop(col_info["description"]) embed_columns.extend(embedding_result) idx = 0 diff --git a/api/loaders/mysql_loader.py b/api/loaders/mysql_loader.py index fc922ac3..92f0af58 100644 --- a/api/loaders/mysql_loader.py +++ b/api/loaders/mysql_loader.py @@ -1,5 +1,6 @@ """MySQL loader for loading database schemas into FalkorDB graphs.""" +import asyncio import datetime import decimal import logging @@ -173,8 +174,17 @@ async def load( # pylint: disable=arguments-differ # Parse connection URL conn_params = MySQLLoader._parse_mysql_url(connection_url) - # Connect to MySQL database - conn = pymysql.connect(**conn_params) + # Off-loop: pymysql is a blocking driver and this generator backs + # the connect/refresh streaming responses. Inline introspection + # blocks the event loop, so those streams cannot emit keepalives + # while a large schema loads. Only the connect is time-bounded + # here; introspecting a very large schema can legitimately outlast + # DB_STATEMENT_TIMEOUT, which is sized for user queries. + conn = await asyncio.to_thread( + pymysql.connect, + connect_timeout=Config.DB_CONNECT_TIMEOUT, + **conn_params, + ) cursor = conn.cursor(DictCursor) # Get database name @@ -182,11 +192,15 @@ async def load( # pylint: disable=arguments-differ # Get all table information yield True, "Extracting table information..." - entities = MySQLLoader.extract_tables_info(cursor, db_name) + entities = await asyncio.to_thread( + MySQLLoader.extract_tables_info, cursor, db_name + ) # Get all relationship information yield True, "Extracting relationship information..." - relationships = MySQLLoader.extract_relationships(cursor, db_name) + relationships = await asyncio.to_thread( + MySQLLoader.extract_relationships, cursor, db_name + ) # Close database connection cursor.close() diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index b0325689..6f9ca420 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -1,5 +1,6 @@ """PostgreSQL loader for loading database schemas into FalkorDB graphs.""" +import asyncio import re import datetime import decimal @@ -165,14 +166,25 @@ async def load( # pylint: disable=arguments-differ # Parse schema from connection URL (defaults to 'public') schema = PostgresLoader.parse_schema_from_url(connection_url) - # Connect to PostgreSQL database - conn = psycopg2.connect(connection_url) + # Off-loop: psycopg2 is a blocking driver and this generator backs + # the connect/refresh streaming responses. Running introspection + # inline blocks the event loop, so those streams cannot emit + # keepalives while a large schema loads. Only the connect is + # time-bounded here: introspecting a very large schema can + # legitimately outlast DB_STATEMENT_TIMEOUT, which is sized for + # user queries. + conn = await asyncio.to_thread( + psycopg2.connect, + connection_url, + connect_timeout=Config.DB_CONNECT_TIMEOUT, + ) cursor = conn.cursor() # Set the session search_path to the parsed schema so unqualified # table references (e.g. in sample queries) resolve correctly. - cursor.execute( - sql.SQL("SET search_path TO {}").format(sql.Identifier(schema)) + await asyncio.to_thread( + cursor.execute, + sql.SQL("SET search_path TO {}").format(sql.Identifier(schema)), ) # Extract database name from connection URL @@ -182,11 +194,15 @@ async def load( # pylint: disable=arguments-differ # Get all table information yield True, "Extracting table information..." - entities = PostgresLoader.extract_tables_info(cursor, schema) + entities = await asyncio.to_thread( + PostgresLoader.extract_tables_info, cursor, schema + ) yield True, "Extracting relationship information..." # Get all relationship information - relationships = PostgresLoader.extract_relationships(cursor, schema) + relationships = await asyncio.to_thread( + PostgresLoader.extract_relationships, cursor, schema + ) # Close database connection before graph loading cursor.close() @@ -553,22 +569,40 @@ def _execution_connect_kwargs(db_url: str) -> Dict[str, Any]: url_options = url_params.get("options", [""])[0] kwargs: Dict[str, Any] = {} + # The configured values are maximums, not defaults: a URL may tighten + # them but must not loosen them or switch them off. In libpq and + # psycopg2 a timeout of 0 means "no limit", which would let one query + # hold a shared worker thread indefinitely. + # # Match an actual directive, not the bare word: a substring test would # also hit something like ``-c application_name=statement_timeout_probe`` # and silently skip our bound. - has_statement_timeout = re.search( - r"(?:^|\s)-c\s*statement_timeout\s*=", url_options + timeout_ms = Config.DB_STATEMENT_TIMEOUT * 1000 + url_statement_timeout = re.search( + r"(?:^|\s)-c\s*statement_timeout\s*=\s*(\d+)", url_options ) - if has_statement_timeout: + if url_statement_timeout and 0 < int(url_statement_timeout.group(1)) <= timeout_ms: options = url_options else: - timeout_ms = Config.DB_STATEMENT_TIMEOUT * 1000 - options = f"{url_options} -c statement_timeout={timeout_ms}".strip() + # Drop any existing directive so ours is not merely appended + # alongside a disabled or looser value. + stripped = re.sub( + r"(?:^|\s)-c\s*statement_timeout\s*=\s*\S*", " ", url_options + ).strip() + options = f"{stripped} -c statement_timeout={timeout_ms}".strip() if options: kwargs["options"] = options - if "connect_timeout" not in url_params: - kwargs["connect_timeout"] = Config.DB_CONNECT_TIMEOUT + url_connect_timeout = url_params.get("connect_timeout", [None])[0] + connect_timeout = Config.DB_CONNECT_TIMEOUT + if url_connect_timeout is not None: + try: + requested = int(url_connect_timeout) + except ValueError: + requested = 0 + if 0 < requested <= connect_timeout: + connect_timeout = requested + kwargs["connect_timeout"] = connect_timeout return kwargs @staticmethod diff --git a/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index 4dd7d246..e6f3af7c 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -1,5 +1,6 @@ """Snowflake loader for loading database schemas into FalkorDB graphs.""" +import asyncio import base64 import datetime import decimal @@ -259,8 +260,13 @@ async def load(prefix: str, connection_url: str) -> AsyncGenerator[ # Parse connection URL conn_params = SnowflakeLoader._parse_snowflake_url(connection_url) - # Connect to Snowflake database - conn = snowflake.connector.connect(**conn_params) + # Off-loop, as in the other loaders: this generator backs the + # connect/refresh streaming responses, and inline introspection + # blocks the event loop so no stream can emit keepalives. The + # parser's login/network timeouts already bound the connect. + conn = await asyncio.to_thread( + snowflake.connector.connect, **conn_params + ) cursor = conn.cursor(DictCursor) # Get database and schema name @@ -271,11 +277,15 @@ async def load(prefix: str, connection_url: str) -> AsyncGenerator[ # Get all table information yield True, "Extracting table information..." - entities = SnowflakeLoader.extract_tables_info(cursor, db_name, schema_name) + entities = await asyncio.to_thread( + SnowflakeLoader.extract_tables_info, cursor, db_name, schema_name + ) # Get all relationship information yield True, "Extracting relationship information..." - relationships = SnowflakeLoader.extract_relationships(cursor, db_name, schema_name) + relationships = await asyncio.to_thread( + SnowflakeLoader.extract_relationships, cursor, db_name, schema_name + ) # Close database connection cursor.close() diff --git a/api/memory/graphiti_tool.py b/api/memory/graphiti_tool.py index f40551ec..eb4e407c 100644 --- a/api/memory/graphiti_tool.py +++ b/api/memory/graphiti_tool.py @@ -28,6 +28,10 @@ from api.agents.utils import run_completion +from api.embeddings import ( + embed_one_off_loop, + vector_size_off_loop, +) def extract_embedding_model_name(full_model_name: str) -> str: @@ -96,7 +100,7 @@ async def create( await self._ensure_entity_nodes_direct(user_id, graph_id) - vector_size = Config.EMBEDDING_MODEL.get_vector_size() + vector_size = await vector_size_off_loop() driver = self.graphiti_client.driver await driver.execute_query(f"CREATE VECTOR INDEX FOR (p:Query) ON (p.embeddings) OPTIONS {{dimension:{vector_size}, similarityFunction:'euclidean'}}") @@ -124,7 +128,7 @@ async def _ensure_entity_nodes_direct(self, user_id: str, database_name: str) -> if not user_check_result[0]: # If no records found, create user node user_uuid = str(uuid.uuid4()) - user_name_embedding = Config.EMBEDDING_MODEL.embed(user_node_name)[0] + user_name_embedding = await embed_one_off_loop(user_node_name) user_node_data = { 'uuid': user_uuid, @@ -161,7 +165,7 @@ async def _ensure_entity_nodes_direct(self, user_id: str, database_name: str) -> if not database_check_result[0]: # If no records found, create database node database_uuid = str(uuid.uuid4()) - database_name_embedding = Config.EMBEDDING_MODEL.embed(database_node_name)[0] + database_name_embedding = await embed_one_off_loop(database_node_name) database_node_data = { 'uuid': database_uuid, @@ -355,7 +359,7 @@ async def save_query_memory(self, query: str, sql_query: str, success: bool, err escaped_query = query.replace("'", "\\'").replace('"', '\\"') escaped_sql = sql_query.replace("'", "\\'").replace('"', '\\"') escaped_error = error.replace("'", "\\'").replace('"', '\\"') if error else "" - embeddings = Config.EMBEDDING_MODEL.embed(escaped_query)[0] + embeddings = await embed_one_off_loop(escaped_query) # First check if a Query node with the same user_query and sql_query already exists check_query = f""" @@ -438,7 +442,7 @@ async def retrieve_similar_queries(self, query: str, limit: int = 5) -> List[Dic if not database_node_exists: return [] - query_embedding = Config.EMBEDDING_MODEL.embed(query)[0] + query_embedding = await embed_one_off_loop(query) cypher_query = f""" CALL db.idx.vector.queryNodes('Query', 'embeddings', 10, vecf32($embedding)) YIELD node, score diff --git a/tests/test_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py index b6bb406e..751a18b0 100644 --- a/tests/test_db_execution_timeouts.py +++ b/tests/test_db_execution_timeouts.py @@ -41,15 +41,39 @@ def test_postgres_preserves_url_options(): @pytest.mark.unit -def test_postgres_url_timeouts_win(): - """An explicit value in the URL is not overridden.""" +def test_postgres_url_may_tighten_the_bounds(): + """A stricter URL value is honoured.""" kwargs = PostgresLoader._execution_connect_kwargs( f"{PG_URL}?options=-c%20statement_timeout%3D1234" ) assert kwargs["options"] == "-c statement_timeout=1234" kwargs = PostgresLoader._execution_connect_kwargs(f"{PG_URL}?connect_timeout=3") - assert "connect_timeout" not in kwargs + assert kwargs["connect_timeout"] == 3 + + +@pytest.mark.unit +@pytest.mark.parametrize("statement_timeout", ["0", "999999999"]) +def test_postgres_url_cannot_loosen_or_disable_statement_timeout(statement_timeout): + """Configured values are maximums, not defaults. + + ``statement_timeout=0`` means "no limit" in libpq, so honouring it would + let one query hold a shared worker thread indefinitely. + """ + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options=-c%20statement_timeout%3D{statement_timeout}" + ) + # Exact equality: the URL directive is replaced, not appended alongside. + assert kwargs["options"] == f"-c statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}" + + +@pytest.mark.unit +@pytest.mark.parametrize("connect_timeout", ["0", "600"]) +def test_postgres_url_cannot_loosen_or_disable_connect_timeout(connect_timeout): + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?connect_timeout={connect_timeout}" + ) + assert kwargs["connect_timeout"] == Config.DB_CONNECT_TIMEOUT @pytest.mark.unit diff --git a/tests/test_embeddings_offloading.py b/tests/test_embeddings_offloading.py new file mode 100644 index 00000000..a98b6835 --- /dev/null +++ b/tests/test_embeddings_offloading.py @@ -0,0 +1,109 @@ +"""Embedding calls must not run on the event loop. + +``EmbeddingsModel.embed`` is a blocking network call. Called directly from an +async method it blocks the loop, and a blocked loop cannot write keepalives — +so a slow embedding kills open streams regardless of the keepalive wrapper. +The memory path (enabled by default in the browser) and the schema loaders both +embed, so both go through ``api.embeddings``. +""" + +import asyncio +import inspect +import time + +import pytest + +from api.config import Config +from api.embeddings import embed_off_loop, embed_one_off_loop, vector_size_off_loop + +STALL = 0.4 +TICK = 0.02 + + +@pytest.fixture(name="slow_embedding") +def _slow_embedding(monkeypatch): + def slow_embed(text): + time.sleep(STALL) + count = len(text) if isinstance(text, list) else 1 + return [[0.1, 0.2] for _ in range(count)] + + def slow_vector_size(): + time.sleep(STALL) + return 2 + + monkeypatch.setattr(Config.EMBEDDING_MODEL, "embed", slow_embed) + monkeypatch.setattr(Config.EMBEDDING_MODEL, "get_vector_size", slow_vector_size) + + +async def _ticks_during(coro): + """Count event-loop ticks while *coro* runs.""" + ticks = [] + stop = asyncio.Event() + + async def ticker(): + while not stop.is_set(): + ticks.append(time.monotonic()) + await asyncio.sleep(TICK) + + ticker_task = asyncio.ensure_future(ticker()) + result = await coro + stop.set() + await ticker_task + return result, ticks + + +@pytest.mark.unit +async def test_embed_off_loop_keeps_the_loop_responsive(slow_embedding): + vectors, ticks = await _ticks_during(embed_off_loop(["a", "b"])) + assert len(vectors) == 2 + assert len(ticks) > (STALL / TICK) * 0.4, f"event loop starved: {len(ticks)} ticks" + + +@pytest.mark.unit +async def test_embed_one_off_loop_returns_a_single_vector(slow_embedding): + vector, ticks = await _ticks_during(embed_one_off_loop("a")) + assert vector == [0.1, 0.2] + assert len(ticks) > (STALL / TICK) * 0.4, f"event loop starved: {len(ticks)} ticks" + + +@pytest.mark.unit +async def test_vector_size_off_loop_keeps_the_loop_responsive(slow_embedding): + size, ticks = await _ticks_during(vector_size_off_loop()) + assert size == 2 + assert len(ticks) > (STALL / TICK) * 0.4, f"event loop starved: {len(ticks)} ticks" + + +@pytest.mark.unit +def test_async_callers_do_not_embed_inline(): + """Guard the call sites: no bare ``EMBEDDING_MODEL.embed`` in async modules. + + These modules run inside streaming responses, so an inline embed there + reintroduces the stall this suite exists to prevent. + """ + import api.graph + import api.loaders.graph_loader as graph_loader + import api.memory.graphiti_tool as graphiti_tool + + offenders = [] + for module in (api.graph, graph_loader, graphiti_tool): + source = inspect.getsource(module) + for lineno, line in enumerate(source.splitlines(), start=1): + stripped = line.strip() + if stripped.startswith("#"): + continue + if "EMBEDDING_MODEL.embed(" in stripped or ( + "EMBEDDING_MODEL.get_vector_size(" in stripped + ): + if "to_thread" not in stripped: + offenders.append(f"{module.__name__}:{lineno}: {stripped}") + + assert not offenders, "inline embedding call(s):\n" + "\n".join(offenders) + + +@pytest.mark.unit +def test_embedding_calls_are_time_bounded(): + """A hung provider must not pin a worker thread forever.""" + kwargs = Config.EMBEDDING_MODEL._embedding_kwargs() + assert kwargs["timeout"] == Config.LLM_TIMEOUT + assert kwargs["max_retries"] == Config.LLM_MAX_RETRIES + assert kwargs["num_retries"] == 0 diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py new file mode 100644 index 00000000..fd161b7a --- /dev/null +++ b/tests/test_schema_load_offloading.py @@ -0,0 +1,98 @@ +"""Schema loading must not block the event loop. + +``load()`` backs the connect and refresh streaming responses. Its driver work — +connect, introspection — is synchronous, and a blocked loop cannot write +keepalives, so inline introspection means those two streams go silent for the +whole load and can be severed by an idle timeout. +""" + +import asyncio +import time +from unittest.mock import MagicMock, patch + +import pytest + +import api.core # noqa: F401 pylint: disable=unused-import +from api.core.pipeline import MySQLLoader, PostgresLoader + +STALL = 0.3 +TICK = 0.02 + + +async def _ticks_while_consuming(agen): + """Drain *agen*, counting event-loop ticks.""" + ticks = [] + stop = asyncio.Event() + + async def ticker(): + while not stop.is_set(): + ticks.append(time.monotonic()) + await asyncio.sleep(TICK) + + ticker_task = asyncio.ensure_future(ticker()) + steps = [step async for step in agen] + stop.set() + await ticker_task + return steps, ticks + + +def _slow(*_args, **_kwargs): + time.sleep(STALL) + return {} + + +@pytest.mark.unit +@patch("api.loaders.postgres_loader.load_to_graph") +@patch("api.loaders.postgres_loader.PostgresLoader.extract_relationships", _slow) +@patch("api.loaders.postgres_loader.PostgresLoader.extract_tables_info", _slow) +@patch("api.loaders.postgres_loader.psycopg2.connect") +async def test_postgres_load_does_not_block_the_loop(mock_connect, mock_load_to_graph): + def slow_connect(*_args, **_kwargs): + time.sleep(STALL) + return MagicMock() + + mock_connect.side_effect = slow_connect + + async def noop(*_args, **_kwargs): + return None + + mock_load_to_graph.side_effect = noop + + _steps, ticks = await _ticks_while_consuming( + PostgresLoader.load("pfx", "postgresql://u:p@h:5432/db") + ) + + # Three blocking stages at STALL each; the loop must stay responsive. + assert len(ticks) > (STALL * 3 / TICK) * 0.3, ( + f"event loop starved during schema load: {len(ticks)} ticks" + ) + gaps = [b - a for a, b in zip(ticks, ticks[1:])] + assert max(gaps) < STALL, f"loop blocked for {max(gaps):.2f}s" + + +@pytest.mark.unit +@patch("api.loaders.mysql_loader.load_to_graph") +@patch("api.loaders.mysql_loader.MySQLLoader.extract_relationships", _slow) +@patch("api.loaders.mysql_loader.MySQLLoader.extract_tables_info", _slow) +@patch("api.loaders.mysql_loader.pymysql.connect") +async def test_mysql_load_does_not_block_the_loop(mock_connect, mock_load_to_graph): + def slow_connect(*_args, **_kwargs): + time.sleep(STALL) + return MagicMock() + + mock_connect.side_effect = slow_connect + + async def noop(*_args, **_kwargs): + return None + + mock_load_to_graph.side_effect = noop + + _steps, ticks = await _ticks_while_consuming( + MySQLLoader.load("pfx", "mysql://u:p@h:3306/db") + ) + + assert len(ticks) > (STALL * 3 / TICK) * 0.3, ( + f"event loop starved during schema load: {len(ticks)} ticks" + ) + gaps = [b - a for a, b in zip(ticks, ticks[1:])] + assert max(gaps) < STALL, f"loop blocked for {max(gaps):.2f}s" diff --git a/tests/test_stream_idle_timeout.py b/tests/test_stream_idle_timeout.py index b5c8ee69..3ae58b62 100644 --- a/tests/test_stream_idle_timeout.py +++ b/tests/test_stream_idle_timeout.py @@ -194,3 +194,56 @@ async def test_slow_sql_execution_keeps_stream_alive(pipeline_stubs): assert keepalives >= 2, "no keepalive during the SQL-execution stall" assert max_gap < MAX_IDLE, f"stream idle for {max_gap:.2f}s during SQL execution" assert any('"query_result"' in p for p in payloads) + + +@pytest.mark.unit +async def test_slow_memory_search_keeps_stream_alive(pipeline_stubs, monkeypatch): + """Stage 4: memory search, which the browser enables by default. + + ``ChatInterface`` sends ``useMemory = true``, so this path is on for real + traffic even though the request model defaults it to ``False``. The lookup + runs inside the silent window before the SQL chunk, and it embeds the query + — a blocking network call. + """ + t2s = pipeline_stubs["text2sql"] + + class _MemoryTool: + async def search_memories(self, query): + # Mirrors the real path: the embedding is offloaded, not inline. + await asyncio.to_thread(time.sleep, STALL) + return "previously asked about accounts" + + async def fake_create_memory_tool(user_id, graph_id, db=None): + return _MemoryTool() + + monkeypatch.setattr(t2s, "_create_memory_tool", fake_create_memory_tool) + + class _MemoryChat(_Chat): + def __init__(self): + super().__init__() + self.use_memory = True + + from api.core.text2sql import _Final + + async def serialize(gen): + async for event in gen: + if isinstance(event, _Final): + break + yield json.dumps(event) + MESSAGE_DELIMITER + + gaps, keepalives, payloads = [], 0, [] + last = time.monotonic() + async for chunk in with_keepalive( + serialize(t2s.run_query("u", "g", _MemoryChat())), interval=INTERVAL + ): + now = time.monotonic() + gaps.append(now - last) + last = now + if chunk == MESSAGE_DELIMITER: + keepalives += 1 + else: + payloads.append(chunk) + + assert keepalives >= 2, "no keepalive during the memory-search stall" + assert max(gaps) < MAX_IDLE, f"stream idle for {max(gaps):.2f}s during memory search" + assert any('"ai_response"' in p for p in payloads) From 5da0770d3c748397f5a05d8fba4c331098be97e5 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Thu, 20 Aug 2026 10:02:26 +0300 Subject: [PATCH 10/25] test: address lint-bot nits in the new offloading tests - Drop the redundant `import api.core` from the schema-load test: importing `api.core.pipeline` already initialises the package, so the side-effect import was unnecessary here (unlike in the execution-timeout test, where the loaders are reached directly). - Have the ticker helpers return their samples instead of appending to a closure and ending on a bare `await task` statement, which the scanner flags as having no effect. Same pattern in all three offloading tests. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_embeddings_offloading.py | 8 ++++---- tests/test_find_offloading.py | 7 ++++--- tests/test_schema_load_offloading.py | 9 ++++----- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/test_embeddings_offloading.py b/tests/test_embeddings_offloading.py index a98b6835..2ede3a34 100644 --- a/tests/test_embeddings_offloading.py +++ b/tests/test_embeddings_offloading.py @@ -37,19 +37,19 @@ def slow_vector_size(): async def _ticks_during(coro): """Count event-loop ticks while *coro* runs.""" - ticks = [] stop = asyncio.Event() async def ticker(): + samples = [] while not stop.is_set(): - ticks.append(time.monotonic()) + samples.append(time.monotonic()) await asyncio.sleep(TICK) + return samples ticker_task = asyncio.ensure_future(ticker()) result = await coro stop.set() - await ticker_task - return result, ticks + return result, await ticker_task @pytest.mark.unit diff --git a/tests/test_find_offloading.py b/tests/test_find_offloading.py index c6d369dd..5ba1b888 100644 --- a/tests/test_find_offloading.py +++ b/tests/test_find_offloading.py @@ -66,20 +66,21 @@ def slow_embed(texts): @pytest.mark.unit async def test_find_does_not_block_the_event_loop(slow_find_deps): """A ticker must keep running while find() is in its blocking calls.""" - ticks = [] stop = asyncio.Event() async def ticker(): + samples = [] while not stop.is_set(): - ticks.append(time.monotonic()) + samples.append(time.monotonic()) await asyncio.sleep(TICK) + return samples ticker_task = asyncio.ensure_future(ticker()) started = time.monotonic() tables = await graph_module.find("g", ["show me customers"], "CRM demo.") elapsed = time.monotonic() - started stop.set() - await ticker_task + ticks = await ticker_task # Both blocking calls ran, so this took at least 2 * STALL. assert elapsed >= STALL * 2 * 0.9, f"stalls did not run (elapsed {elapsed:.2f}s)" diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py index fd161b7a..d5c17f01 100644 --- a/tests/test_schema_load_offloading.py +++ b/tests/test_schema_load_offloading.py @@ -12,7 +12,6 @@ import pytest -import api.core # noqa: F401 pylint: disable=unused-import from api.core.pipeline import MySQLLoader, PostgresLoader STALL = 0.3 @@ -21,19 +20,19 @@ async def _ticks_while_consuming(agen): """Drain *agen*, counting event-loop ticks.""" - ticks = [] stop = asyncio.Event() async def ticker(): + samples = [] while not stop.is_set(): - ticks.append(time.monotonic()) + samples.append(time.monotonic()) await asyncio.sleep(TICK) + return samples ticker_task = asyncio.ensure_future(ticker()) steps = [step async for step in agen] stop.set() - await ticker_task - return steps, ticks + return steps, await ticker_task def _slow(*_args, **_kwargs): From 88fed2e486a560f86d7e0176bded6ba4397469ea Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Thu, 20 Aug 2026 12:03:28 +0300 Subject: [PATCH 11/25] fix: stop orphaning speculative work, confine DB work to one worker, harden clamps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fourth review from @Naseem77; all four findings were valid, and the first two are consequences of the `to_thread` offloading added earlier in this PR. **1. Off-topic requests orphaned speculative work.** `find_task` and `memory_tool_task` were started before the relevancy check. Cancelling a task whose thread is blocked in a socket read does not stop that thread, so an off-topic question abandoned the task while the provider call ran to completion — consuming executor capacity and provider quota after the response had been sent, and `memory_tool_task` was never cancelled or awaited at all. Repeated off-topic requests could saturate the thread pool every other offloaded call depends on. Relevancy now runs first, and the concurrent work starts only once the question is known to be answerable; the two tasks are gathered together so neither is left unobserved if the other fails. Cost is one relevancy round-trip on answerable questions, which the original code called a "small perf win" in the other direction. Verified on the harness: an off-topic query now logs only `llm_call label=relevancy` — no find, no embedding. **2. Schema loading mishandled resources on cancellation.** PostgreSQL closed the cursor and connection from the generator's `finally`, which can run while an offloaded introspection is still using them — two threads on one connection. MySQL and Snowflake had no `finally` at all, so any failure or disconnect leaked the session outright. Connect, cursor, introspection and cleanup now live in a single worker (`_introspect_schema`) with `try/finally`, so the thread that owns the resources is the one that closes them. **3. The timeout clamp was bypassable.** libpq applies the last directive, so `statement_timeout=1000 ... statement_timeout=0` ended up unbounded, and a unit-bearing value like `2min` passed the digit check on its leading digits. Every accepted directive is now stripped and exactly one normalised bound appended; a URL value is honoured only when unambiguous — a single directive, plain milliseconds, no looser than the ceiling. **4. Zero timeouts were accepted from the environment.** Zero disables the PostgreSQL limit entirely and makes PyMySQL raise at query time. Timeout config is now validated at import: non-positive or non-numeric values fail fast with a message naming the variable. `LLM_MAX_RETRIES=0` remains valid — it means "no retry", a stricter ceiling. Tests: 18 new cases across `test_config_validation.py` (new), `test_db_execution_timeouts.py` and `test_schema_load_offloading.py`, covering the five bypass shapes, cancellation cleanup, and cleanup on a failed introspection. The last has teeth: neutering the worker's `finally` fails with "connection leaked when introspection failed". 264 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/config.py | 37 ++++++-- api/core/text2sql.py | 50 ++++++----- api/loaders/mysql_loader.py | 60 +++++++------ api/loaders/postgres_loader.py | 126 ++++++++++++++------------- api/loaders/snowflake_loader.py | 57 +++++++----- tests/test_db_execution_timeouts.py | 34 ++++++++ tests/test_schema_load_offloading.py | 86 ++++++++++++++++++ 7 files changed, 310 insertions(+), 140 deletions(-) diff --git a/api/config.py b/api/config.py index 9f83b975..63e645f9 100644 --- a/api/config.py +++ b/api/config.py @@ -112,6 +112,29 @@ def _with_prefix(model: str, provider: str) -> str: SUPPORTED_VENDORS = ("openai", "anthropic", "gemini", "azure", "ollama", "cohere") +def _positive_env(name: str, default: str, cast=int): + """Read a timeout-style env var, rejecting values that disable the bound. + + Zero is not a harmless "unset" here: PostgreSQL treats a 0 timeout as + "no limit", which removes the safeguard entirely, and PyMySQL raises at + query time on a 0 socket timeout. Fail at startup with a clear message + rather than silently losing the protection. + """ + raw = os.getenv(name, default) + try: + value = cast(raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"{name} must be a positive number (got {raw!r})" + ) from exc + if value <= 0: + raise ValueError( + f"{name} must be greater than 0 (got {raw!r}); a zero or negative " + "timeout disables the safeguard it exists to provide" + ) + return value + + @dataclasses.dataclass class Config: """ @@ -172,13 +195,15 @@ class Config: # through to litellm, which aborts the underlying HTTP request — so a # hung provider surfaces as a clean error instead of stalling the # response stream indefinitely (incident 2026-07-29). - LLM_TIMEOUT: float = float(os.getenv("LLM_TIMEOUT", "90")) # pylint: disable=invalid-name + LLM_TIMEOUT: float = _positive_env("LLM_TIMEOUT", "90", float) # pylint: disable=invalid-name # A call slower than this is logged at WARNING. Normal analysis calls # completed in ~6s during the incident window, so this flags outliers # well before they reach the timeout. # pylint: disable-next=invalid-name - LLM_SLOW_CALL_THRESHOLD: float = float(os.getenv("LLM_SLOW_CALL_THRESHOLD", "20")) + LLM_SLOW_CALL_THRESHOLD: float = _positive_env( + "LLM_SLOW_CALL_THRESHOLD", "20", float + ) # Retry budget for a single agent LLM call. Kept explicit because the # provider SDK and litellm each have their own retry loop, and leaving @@ -186,16 +211,18 @@ class Config: # 3s timeout took 10.8s to fail). Applied as the SDK-level retry count # with litellm's outer loop disabled, so the worst case stays close to # LLM_TIMEOUT rather than a multiple of it. + # Zero is valid here (it means "no retry", a strict ceiling); negative is + # not. # pylint: disable-next=invalid-name - LLM_MAX_RETRIES: int = int(os.getenv("LLM_MAX_RETRIES", "1")) + LLM_MAX_RETRIES: int = max(0, int(os.getenv("LLM_MAX_RETRIES", "1"))) # Bounds for user-query execution against the target database. Offloading # execution to a thread stops a slow query from blocking other requests, # but nothing bounds how long the query itself runs without these. # pylint: disable-next=invalid-name - DB_CONNECT_TIMEOUT: int = int(os.getenv("DB_CONNECT_TIMEOUT", "10")) + DB_CONNECT_TIMEOUT: int = _positive_env("DB_CONNECT_TIMEOUT", "10") # pylint: disable-next=invalid-name - DB_STATEMENT_TIMEOUT: int = int(os.getenv("DB_STATEMENT_TIMEOUT", "60")) + DB_STATEMENT_TIMEOUT: int = _positive_env("DB_STATEMENT_TIMEOUT", "60") DB_MAX_DISTINCT: int = 100 # pylint: disable=invalid-name DB_UNIQUENESS_THRESHOLD: float = 0.5 # pylint: disable=invalid-name diff --git a/api/core/text2sql.py b/api/core/text2sql.py index ba0d4c8c..b50743db 100644 --- a/api/core/text2sql.py +++ b/api/core/text2sql.py @@ -338,13 +338,6 @@ async def run_query( # pylint: disable=too-many-locals,too-many-branches,too-ma logging.info("User Query: %s", sanitize_query(queries_history[-1])) - # Memory tool created concurrently with relevancy/find work — small perf - # win for streaming, harmless for SDK. Lazy-imported via _create_memory_tool. - memory_tool_task = ( - asyncio.create_task(_create_memory_tool(user_id, namespaced, db=db)) - if use_memory else None - ) - yield { "type": "reasoning_step", "final_response": False, @@ -369,24 +362,21 @@ async def run_query( # pylint: disable=too-many-locals,too-many-branches,too-ma )) return - # Concurrent: relevancy check + table-finding - find_task = asyncio.create_task( - find(namespaced, queries_history, db_description, db=db) - ) + # Relevancy runs before table-finding and memory-tool creation, not + # alongside them. Both of those make provider calls in worker threads, and + # a thread blocked in a socket read cannot be cancelled from Python: on an + # off-topic question, cancelling the task abandons the *task* while the + # call keeps running to completion, consuming executor capacity and + # provider quota after the response has already been sent. Repeated + # off-topic requests could saturate the thread pool that every other + # offloaded call depends on. Sequencing costs one relevancy round-trip on + # answerable questions and starts no work that cannot be used. agent_rel = RelevancyAgent( queries_history, result_history, custom_api_key, custom_model, ) - relevancy_task = asyncio.create_task( - agent_rel.get_answer(queries_history[-1], db_description) - ) - answer_rel = await relevancy_task + answer_rel = await agent_rel.get_answer(queries_history[-1], db_description) if answer_rel["status"] != "On-topic": - find_task.cancel() - try: - await find_task - except asyncio.CancelledError: - logging.debug("Find task cancelled (off-topic query)") msg = "Off topic question: " + answer_rel["reason"] yield {"type": "followup_questions", "final_response": True, "message": msg} yield _Final(_build_query_result( @@ -396,12 +386,28 @@ async def run_query( # pylint: disable=too-many-locals,too-many-branches,too-ma )) return - tables = await find_task + # Concurrent now that both results are certain to be used. Gathered + # together so neither task is left unobserved if the other fails. + find_task = asyncio.create_task( + find(namespaced, queries_history, db_description, db=db) + ) + memory_tool_task = ( + asyncio.create_task(_create_memory_tool(user_id, namespaced, db=db)) + if use_memory else None + ) + pending = [t for t in (find_task, memory_tool_task) if t is not None] + gathered = await asyncio.gather(*pending, return_exceptions=True) + + tables = gathered[0] + if isinstance(tables, BaseException): + raise tables memory_tool = None memory_context = None if memory_tool_task is not None: - memory_tool = await memory_tool_task + memory_tool = gathered[1] + if isinstance(memory_tool, BaseException): + raise memory_tool memory_context = await memory_tool.search_memories(query=queries_history[-1]) agent_an = AnalysisAgent( diff --git a/api/loaders/mysql_loader.py b/api/loaders/mysql_loader.py index 92f0af58..e81d82c6 100644 --- a/api/loaders/mysql_loader.py +++ b/api/loaders/mysql_loader.py @@ -153,6 +153,36 @@ def _parse_mysql_url(connection_url: str) -> Dict[str, str]: 'database': database } + @staticmethod + def _introspect_schema(conn_params: Dict[str, Any], db_name: str): + """Connect, introspect and close — all inside one worker thread. + + Cleanup lives in the same thread that owns the connection. Cancelling a + ``to_thread`` call does not stop the thread, so closing from the event + loop could run alongside an in-flight introspection; and without a + ``finally`` here a client disconnect leaked the connection outright, + which exhausts database sessions under repeated disconnects. + + Only the connect is time-bounded: introspecting a very large schema can + legitimately outlast DB_STATEMENT_TIMEOUT, which is sized for user + queries. + """ + conn = None + cursor = None + try: + conn = pymysql.connect( + connect_timeout=Config.DB_CONNECT_TIMEOUT, **conn_params + ) + cursor = conn.cursor(DictCursor) + entities = MySQLLoader.extract_tables_info(cursor, db_name) + relationships = MySQLLoader.extract_relationships(cursor, db_name) + return entities, relationships + finally: + if cursor is not None: + cursor.close() + if conn is not None: + conn.close() + @staticmethod async def load( # pylint: disable=arguments-differ prefix: str, @@ -173,39 +203,13 @@ async def load( # pylint: disable=arguments-differ try: # Parse connection URL conn_params = MySQLLoader._parse_mysql_url(connection_url) - - # Off-loop: pymysql is a blocking driver and this generator backs - # the connect/refresh streaming responses. Inline introspection - # blocks the event loop, so those streams cannot emit keepalives - # while a large schema loads. Only the connect is time-bounded - # here; introspecting a very large schema can legitimately outlast - # DB_STATEMENT_TIMEOUT, which is sized for user queries. - conn = await asyncio.to_thread( - pymysql.connect, - connect_timeout=Config.DB_CONNECT_TIMEOUT, - **conn_params, - ) - cursor = conn.cursor(DictCursor) - - # Get database name db_name = conn_params['database'] - # Get all table information yield True, "Extracting table information..." - entities = await asyncio.to_thread( - MySQLLoader.extract_tables_info, cursor, db_name - ) - - # Get all relationship information - yield True, "Extracting relationship information..." - relationships = await asyncio.to_thread( - MySQLLoader.extract_relationships, cursor, db_name + entities, relationships = await asyncio.to_thread( + MySQLLoader._introspect_schema, conn_params, db_name ) - # Close database connection - cursor.close() - conn.close() - # Load data into graph yield True, "Loading data into graph..." await load_to_graph(f"{prefix}_{db_name}", entities, relationships, diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index 6f9ca420..1239bdc9 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -141,6 +141,45 @@ def parse_schema_from_url(connection_url: str) -> str: except Exception: # pylint: disable=broad-exception-caught return 'public' + @staticmethod + def _introspect_schema(connection_url: str, schema: str): + """Connect, introspect and close — all inside one worker thread. + + Everything touching the driver lives here so the connection and cursor + are created, used and closed by the same thread. Closing them from the + event loop instead (in the generator's ``finally``) can run while an + offloaded introspection is still using them, because cancelling a + ``to_thread`` call does not stop the thread it is running in: the + result is two threads on one connection. Keeping cleanup in the worker + also means a client disconnect cannot leak the connection. + + Only the connect is time-bounded: introspecting a very large schema can + legitimately outlast DB_STATEMENT_TIMEOUT, which is sized for user + queries. + """ + conn = None + cursor = None + try: + conn = psycopg2.connect( + connection_url, connect_timeout=Config.DB_CONNECT_TIMEOUT + ) + cursor = conn.cursor() + + # Set the session search_path to the parsed schema so unqualified + # table references (e.g. in sample queries) resolve correctly. + cursor.execute( + sql.SQL("SET search_path TO {}").format(sql.Identifier(schema)) + ) + + entities = PostgresLoader.extract_tables_info(cursor, schema) + relationships = PostgresLoader.extract_relationships(cursor, schema) + return entities, relationships + finally: + if cursor is not None: + cursor.close() + if conn is not None: + conn.close() + @staticmethod async def load( # pylint: disable=arguments-differ prefix: str, @@ -160,56 +199,20 @@ async def load( # pylint: disable=arguments-differ Returns: Tuple[bool, str]: Success status and message """ - conn = None - cursor = None try: # Parse schema from connection URL (defaults to 'public') schema = PostgresLoader.parse_schema_from_url(connection_url) - # Off-loop: psycopg2 is a blocking driver and this generator backs - # the connect/refresh streaming responses. Running introspection - # inline blocks the event loop, so those streams cannot emit - # keepalives while a large schema loads. Only the connect is - # time-bounded here: introspecting a very large schema can - # legitimately outlast DB_STATEMENT_TIMEOUT, which is sized for - # user queries. - conn = await asyncio.to_thread( - psycopg2.connect, - connection_url, - connect_timeout=Config.DB_CONNECT_TIMEOUT, - ) - cursor = conn.cursor() - - # Set the session search_path to the parsed schema so unqualified - # table references (e.g. in sample queries) resolve correctly. - await asyncio.to_thread( - cursor.execute, - sql.SQL("SET search_path TO {}").format(sql.Identifier(schema)), - ) - # Extract database name from connection URL db_name = connection_url.split('/')[-1] if '?' in db_name: db_name = db_name.split('?')[0] - # Get all table information yield True, "Extracting table information..." - entities = await asyncio.to_thread( - PostgresLoader.extract_tables_info, cursor, schema - ) - - yield True, "Extracting relationship information..." - # Get all relationship information - relationships = await asyncio.to_thread( - PostgresLoader.extract_relationships, cursor, schema + entities, relationships = await asyncio.to_thread( + PostgresLoader._introspect_schema, connection_url, schema ) - # Close database connection before graph loading - cursor.close() - cursor = None - conn.close() - conn = None - yield True, "Loading data into graph..." # Load data into graph await load_to_graph(f"{prefix}_{db_name}", entities, relationships, @@ -224,11 +227,6 @@ async def load( # pylint: disable=arguments-differ except Exception as e: # pylint: disable=broad-exception-caught logging.error("Error loading PostgreSQL schema: %s", e) yield False, "Failed to load PostgreSQL database schema" - finally: - if cursor is not None: - cursor.close() - if conn is not None: - conn.close() @staticmethod def extract_tables_info(cursor: Any, schema: str = 'public') -> Dict[str, Any]: @@ -570,26 +568,32 @@ def _execution_connect_kwargs(db_url: str) -> Dict[str, Any]: kwargs: Dict[str, Any] = {} # The configured values are maximums, not defaults: a URL may tighten - # them but must not loosen them or switch them off. In libpq and - # psycopg2 a timeout of 0 means "no limit", which would let one query - # hold a shared worker thread indefinitely. + # them but must not loosen them or switch them off. In libpq a timeout + # of 0 means "no limit", which would let one query hold an + # uncancellable worker thread indefinitely. # - # Match an actual directive, not the bare word: a substring test would - # also hit something like ``-c application_name=statement_timeout_probe`` - # and silently skip our bound. + # Every accepted directive is removed and exactly one normalised bound + # appended. Leaving any in place is not safe: libpq applies the last + # occurrence, so `statement_timeout=1000 ... statement_timeout=0` would + # end up unbounded, and a unit-bearing value like `2min` is not + # comparable to our millisecond ceiling. + # + # The pattern matches a real directive, not the bare word: a substring + # test would also hit `-c application_name=statement_timeout_probe`. timeout_ms = Config.DB_STATEMENT_TIMEOUT * 1000 - url_statement_timeout = re.search( - r"(?:^|\s)-c\s*statement_timeout\s*=\s*(\d+)", url_options - ) - if url_statement_timeout and 0 < int(url_statement_timeout.group(1)) <= timeout_ms: - options = url_options - else: - # Drop any existing directive so ours is not merely appended - # alongside a disabled or looser value. - stripped = re.sub( - r"(?:^|\s)-c\s*statement_timeout\s*=\s*\S*", " ", url_options - ).strip() - options = f"{stripped} -c statement_timeout={timeout_ms}".strip() + directive = re.compile(r"(?:^|\s)-c\s*statement_timeout\s*=\s*(\S*)") + found = directive.findall(url_options) + stripped = directive.sub(" ", url_options).strip() + + # Honour a URL value only when it is unambiguous: a single directive, + # plain milliseconds, and no looser than the configured ceiling. + effective_ms = timeout_ms + if len(found) == 1 and found[0].isdigit(): + requested = int(found[0]) + if 0 < requested <= timeout_ms: + effective_ms = requested + + options = f"{stripped} -c statement_timeout={effective_ms}".strip() if options: kwargs["options"] = options diff --git a/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index e6f3af7c..936e04df 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -242,6 +242,36 @@ def _parse_snowflake_url(connection_url: str) -> Dict[str, Any]: # pylint: disa return conn_params + @staticmethod + def _introspect_schema( + conn_params: Dict[str, Any], db_name: str, schema_name: str + ): + """Connect, introspect and close — all inside one worker thread. + + Same reasoning as the other loaders: cleanup belongs in the thread that + owns the connection. Cancelling a ``to_thread`` call does not stop the + thread, and without a ``finally`` here a client disconnect leaked the + session outright. The parser's login/network timeouts bound the + connect. + """ + conn = None + cursor = None + try: + conn = snowflake.connector.connect(**conn_params) + cursor = conn.cursor(DictCursor) + entities = SnowflakeLoader.extract_tables_info( + cursor, db_name, schema_name + ) + relationships = SnowflakeLoader.extract_relationships( + cursor, db_name, schema_name + ) + return entities, relationships + finally: + if cursor is not None: + cursor.close() + if conn is not None: + conn.close() + @staticmethod async def load(prefix: str, connection_url: str) -> AsyncGenerator[ tuple[bool, str], None @@ -259,38 +289,17 @@ async def load(prefix: str, connection_url: str) -> AsyncGenerator[ try: # Parse connection URL conn_params = SnowflakeLoader._parse_snowflake_url(connection_url) - - # Off-loop, as in the other loaders: this generator backs the - # connect/refresh streaming responses, and inline introspection - # blocks the event loop so no stream can emit keepalives. The - # parser's login/network timeouts already bound the connect. - conn = await asyncio.to_thread( - snowflake.connector.connect, **conn_params - ) - cursor = conn.cursor(DictCursor) - - # Get database and schema name db_name = conn_params['database'] # Snowflake stores unquoted identifiers in UPPERCASE; # INFORMATION_SCHEMA lookups require the canonical form. schema_name = conn_params['schema'].upper() - # Get all table information yield True, "Extracting table information..." - entities = await asyncio.to_thread( - SnowflakeLoader.extract_tables_info, cursor, db_name, schema_name - ) - - # Get all relationship information - yield True, "Extracting relationship information..." - relationships = await asyncio.to_thread( - SnowflakeLoader.extract_relationships, cursor, db_name, schema_name + entities, relationships = await asyncio.to_thread( + SnowflakeLoader._introspect_schema, + conn_params, db_name, schema_name, ) - # Close database connection - cursor.close() - conn.close() - # Load data into graph yield True, "Loading data into graph..." await load_to_graph(f"{prefix}_{db_name}", entities, relationships, diff --git a/tests/test_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py index 751a18b0..3445fa53 100644 --- a/tests/test_db_execution_timeouts.py +++ b/tests/test_db_execution_timeouts.py @@ -130,3 +130,37 @@ def test_snowflake_overrides_parser_timeout_defaults(mock_connect): params["session_parameters"]["STATEMENT_TIMEOUT_IN_SECONDS"] == Config.DB_STATEMENT_TIMEOUT ) + + +@pytest.mark.unit +@pytest.mark.parametrize("url_options,reason", [ + ("-c%20statement_timeout%3D1000%20-c%20statement_timeout%3D0", "duplicate, last disables"), + ("-c%20statement_timeout%3D0%20-c%20statement_timeout%3D1000", "duplicate, first disables"), + ("-c%20statement_timeout%3D2min", "unit-bearing value"), + ("-c%20statement_timeout%3D%20", "empty value"), + ("-c%20statement_timeout%3D-5", "negative value"), +]) +def test_postgres_clamp_is_not_bypassable(url_options, reason): + """libpq applies the last directive, so none may survive. + + Leaving an accepted directive in place lets ``statement_timeout=1000 ... + statement_timeout=0`` end up unbounded, and a unit-bearing value like + ``2min`` is not comparable to the millisecond ceiling. + """ + kwargs = PostgresLoader._execution_connect_kwargs(f"{PG_URL}?options={url_options}") + expected = f"-c statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}" + assert kwargs["options"] == expected, reason + assert kwargs["options"].count("statement_timeout") == 1, reason + + +@pytest.mark.unit +def test_postgres_clamp_keeps_unrelated_options(): + """Stripping the timeout directives must not drop other settings.""" + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options=-c%20search_path%3Dfoo%20-c%20statement_timeout%3D0" + ) + assert "search_path=foo" in kwargs["options"] + assert kwargs["options"].endswith( + f"-c statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}" + ) + diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py index d5c17f01..93be18f7 100644 --- a/tests/test_schema_load_offloading.py +++ b/tests/test_schema_load_offloading.py @@ -95,3 +95,89 @@ async def noop(*_args, **_kwargs): ) gaps = [b - a for a, b in zip(ticks, ticks[1:])] assert max(gaps) < STALL, f"loop blocked for {max(gaps):.2f}s" + + +@pytest.mark.unit +@pytest.mark.parametrize("loader_module,loader,url", [ + ("api.loaders.postgres_loader", "PostgresLoader", "postgresql://u:p@h:5432/db"), + ("api.loaders.mysql_loader", "MySQLLoader", "mysql://u:p@h:3306/db"), +]) +async def test_cancelling_a_schema_load_still_closes_the_connection( + loader_module, loader, url +): + """A client disconnect must not leak the database session. + + The introspection runs in a worker thread that cancellation cannot stop, + so cleanup has to live in that same thread. Closing from the generator's + ``finally`` instead would either race the in-flight introspection or, where + there was no ``finally`` at all, leak the connection outright. + """ + conn = MagicMock() + loader_cls = {"PostgresLoader": PostgresLoader, "MySQLLoader": MySQLLoader}[loader] + connect_name = ( + "psycopg2.connect" if loader == "PostgresLoader" else "pymysql.connect" + ) + + def slow_extract(*_args, **_kwargs): + time.sleep(STALL * 3) + return {} + + with patch(f"{loader_module}.{connect_name}", return_value=conn), \ + patch.object(loader_cls, "extract_tables_info", slow_extract), \ + patch.object(loader_cls, "extract_relationships", slow_extract), \ + patch(f"{loader_module}.load_to_graph"): + agen = loader_cls.load("pfx", url) + assert await agen.__anext__() == (True, "Extracting table information...") + + consumer = asyncio.ensure_future(agen.__anext__()) + await asyncio.sleep(STALL) # introspection is in flight + consumer.cancel() + try: + await consumer + except asyncio.CancelledError: + pass + await agen.aclose() + + # The worker owns cleanup, so it runs even though the awaiting task was + # cancelled. Give the thread time to finish and close. + for _ in range(50): + if conn.close.called: + break + await asyncio.sleep(0.05) + + assert conn.close.called, "connection was not closed after cancellation" + + +@pytest.mark.unit +@pytest.mark.parametrize("loader_module,loader,url", [ + ("api.loaders.postgres_loader", "PostgresLoader", "postgresql://u:p@h:5432/db"), + ("api.loaders.mysql_loader", "MySQLLoader", "mysql://u:p@h:3306/db"), +]) +async def test_failed_introspection_still_closes_the_connection( + loader_module, loader, url +): + """An error mid-introspection must not leak the session. + + This is what the worker's ``finally`` buys: MySQL and Snowflake previously + closed only on the success path, so any failure left the connection open, + and repeated failures exhaust database sessions. + """ + conn = MagicMock() + loader_cls = {"PostgresLoader": PostgresLoader, "MySQLLoader": MySQLLoader}[loader] + connect_name = ( + "psycopg2.connect" if loader == "PostgresLoader" else "pymysql.connect" + ) + + def boom(*_args, **_kwargs): + raise RuntimeError("introspection blew up") + + with patch(f"{loader_module}.{connect_name}", return_value=conn), \ + patch.object(loader_cls, "extract_tables_info", boom), \ + patch(f"{loader_module}.load_to_graph"): + steps = [step async for step in loader_cls.load("pfx", url)] + + # The loader reports failure to the stream rather than raising... + assert steps[-1][0] is False + # ...and the connection is closed regardless. + assert conn.close.called, "connection leaked when introspection failed" + From 8e318ff6a0124d003f0128e43272c8c0f4aad40d Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Thu, 20 Aug 2026 12:03:59 +0300 Subject: [PATCH 12/25] test: add the timeout-validation suite that .gitignore silently dropped `tests/test_config_validation.py` matched the repo's `*_conf*` ignore rule, so `git add` skipped it and the previous commit shipped the validation without its tests. Renamed to `test_timeout_validation.py`, which the rule does not match. Covers the four review item #4 cases: zero and negative values rejected for DB_CONNECT_TIMEOUT / DB_STATEMENT_TIMEOUT / LLM_TIMEOUT, non-numeric values rejected, and a clean environment still loading with positive defaults. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_timeout_validation.py | 47 ++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 tests/test_timeout_validation.py diff --git a/tests/test_timeout_validation.py b/tests/test_timeout_validation.py new file mode 100644 index 00000000..99cc0090 --- /dev/null +++ b/tests/test_timeout_validation.py @@ -0,0 +1,47 @@ +"""Timeout configuration must be positive. + +Zero is not a harmless "unset": PostgreSQL treats a 0 timeout as "no limit", +removing the safeguard entirely, and PyMySQL raises at query time on a 0 socket +timeout. Both are worse than refusing to start. +""" + +import importlib + +import pytest + + +def _reload_config(monkeypatch, **env): + for key, value in env.items(): + monkeypatch.setenv(key, value) + import api.config + return importlib.reload(api.config) + + +@pytest.mark.unit +@pytest.mark.parametrize("name", [ + "DB_CONNECT_TIMEOUT", + "DB_STATEMENT_TIMEOUT", + "LLM_TIMEOUT", +]) +@pytest.mark.parametrize("value", ["0", "-1"]) +def test_zero_or_negative_timeouts_are_rejected(monkeypatch, name, value): + with pytest.raises(ValueError, match="greater than 0"): + _reload_config(monkeypatch, **{name: value}) + + +@pytest.mark.unit +def test_non_numeric_timeout_is_rejected(monkeypatch): + with pytest.raises(ValueError, match="positive number"): + _reload_config(monkeypatch, DB_CONNECT_TIMEOUT="abc") + + +@pytest.mark.unit +def test_defaults_are_positive(monkeypatch): + """And a clean environment still loads.""" + for name in ("DB_CONNECT_TIMEOUT", "DB_STATEMENT_TIMEOUT", "LLM_TIMEOUT"): + monkeypatch.delenv(name, raising=False) + module = _reload_config(monkeypatch) + assert module.Config.DB_CONNECT_TIMEOUT > 0 + assert module.Config.DB_STATEMENT_TIMEOUT > 0 + assert module.Config.LLM_TIMEOUT > 0 + assert module.Config.LLM_MAX_RETRIES >= 0 From f565df735100129b1e8a864f4de483bceb33d9a0 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Thu, 20 Aug 2026 12:12:07 +0300 Subject: [PATCH 13/25] test: explain the intentionally empty except in the cancellation test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The scanner flags a bare `pass` with no rationale. The CancelledError is expected — we raised it — and the assertion that matters is the cleanup check after it. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_schema_load_offloading.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py index 93be18f7..8bb0e9c3 100644 --- a/tests/test_schema_load_offloading.py +++ b/tests/test_schema_load_offloading.py @@ -135,6 +135,8 @@ def slow_extract(*_args, **_kwargs): try: await consumer except asyncio.CancelledError: + # Expected: we cancelled it. What matters is the cleanup that runs + # afterwards, asserted below. pass await agen.aclose() From 218834794a30e52e888a223c8829b6fc87a26ec9 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Thu, 20 Aug 2026 12:24:40 +0300 Subject: [PATCH 14/25] fix: offload the last three inline provider calls, and guard against new ones MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Proactive sweep rather than a review response: the same pattern @Naseem77 has flagged four times still existed in three more places, all reachable from a streaming response. - `api/utils.py` `create_combined_description` issues a **batch** completion over every table, and `generate_db_description` a further completion. Both are synchronous and both are called from `load_to_graph`, which backs the connect and refresh streams — so a schema load blocked the event loop for the duration of a batch LLM call over the whole schema. Offloaded at the call sites; `generate_db_description` now goes through `run_completion` for the shared timeout, retry budget and duration logging, and the batch call carries the same bounds. - `api/routes/settings.py` `validate_api_key` called `completion` inline inside an async route, so validating a key against a slow or unreachable provider blocked the loop — and every open query stream — for as long as it took. Offloaded and time-bounded. Behaviour change: that validation call now carries `timeout`, `max_retries=LLM_MAX_RETRIES` and `num_retries=0`. Two tests in `test_settings_route.py` pinned the previous unbounded kwargs and are updated to the bounded contract. The guard in `test_embeddings_offloading.py` is extended from embeddings to all bare provider entry points (`completion(`, `batch_completion(`, `embedding(`) across `api/graph.py`, `graph_loader.py`, `graphiti_tool.py` and `api/routes/settings.py`. Verified with teeth: putting the settings call back inline fails the guard. That check is the part meant to stop this class of bug recurring, rather than finding the next instance by review. Also audited and found clean: the three fire-and-forget task sites in `pipeline.py`, `analytics.py` and `usage_tracking.py` already track their tasks in a sink and attach done-callbacks, so nothing there is unobserved. 265 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/loaders/graph_loader.py | 11 +++++++++-- api/routes/settings.py | 22 ++++++++++++++++----- api/utils.py | 18 +++++++++++------ tests/test_embeddings_offloading.py | 30 +++++++++++++++++++++++++++++ tests/test_settings_route.py | 12 ++++++++++++ 5 files changed, 80 insertions(+), 13 deletions(-) diff --git a/api/loaders/graph_loader.py b/api/loaders/graph_loader.py index 4ff07acc..ce08d3f6 100644 --- a/api/loaders/graph_loader.py +++ b/api/loaders/graph_loader.py @@ -1,5 +1,6 @@ """Graph loader module for loading data into graph databases.""" +import asyncio import json import tqdm @@ -36,7 +37,11 @@ async def load_to_graph( # pylint: disable=too-many-arguments,too-many-position # large schema is loading. vec_len = await vector_size_off_loop() - create_combined_description(entities) + # Both of these make blocking provider calls (a batch completion over every + # table, then a description completion). This coroutine backs the connect + # and refresh streams, so running them inline blocks the event loop and no + # stream can emit keepalives while a schema loads. + await asyncio.to_thread(create_combined_description, entities) try: # Create vector indices @@ -59,7 +64,9 @@ async def load_to_graph( # pylint: disable=too-many-arguments,too-many-position except Exception as e: # pylint: disable=broad-exception-caught print(f"Error creating vector indices: {str(e)}") - db_des = generate_db_description(db_name=db_name, table_names=list(entities.keys())) + db_des = await asyncio.to_thread( + generate_db_description, db_name=db_name, table_names=list(entities.keys()) + ) await graph.query( """ CREATE (d:Database { diff --git a/api/routes/settings.py b/api/routes/settings.py index 554081d2..2e77c49c 100644 --- a/api/routes/settings.py +++ b/api/routes/settings.py @@ -1,11 +1,14 @@ """Settings and configuration routes for the text2sql API.""" +import asyncio +import functools import logging from fastapi import APIRouter, Request from fastapi.responses import JSONResponse from pydantic import BaseModel from litellm import completion +from api.config import Config from api.auth.user_management import token_required from api.routes.tokens import UNAUTHORIZED_RESPONSE @@ -74,11 +77,20 @@ async def validate_api_key(request: Request, data: ValidateKeyRequest): # pylin # Construct model name for LiteLLM (vendor/model format) full_model_name = f"{vendor}/{model}" - test_response = completion( - model=full_model_name, - messages=[{"role": "user", "content": "test"}], - max_tokens=1, - api_key=api_key, + # Off-loop and time-bounded: this is a blocking provider call inside an + # async route, so running it inline blocks the event loop — and with it + # every open query stream — for as long as the provider takes. + test_response = await asyncio.to_thread( + functools.partial( + completion, + model=full_model_name, + messages=[{"role": "user", "content": "test"}], + max_tokens=1, + api_key=api_key, + timeout=Config.LLM_TIMEOUT, + max_retries=Config.LLM_MAX_RETRIES, + num_retries=0, + ) ) # If we get here without exception, the key is valid diff --git a/api/utils.py b/api/utils.py index e6979876..9fd7c6fc 100644 --- a/api/utils.py +++ b/api/utils.py @@ -2,8 +2,9 @@ import json from typing import Dict, List, Optional, TypedDict -from litellm import completion, batch_completion +from litellm import batch_completion +from api.agents.utils import run_completion from api.config import Config @@ -83,11 +84,16 @@ def create_combined_description( # pylint: disable=too-many-locals for batch_start in range(0, len(messages_list), batch_size): batch_messages = messages_list[batch_start : batch_start + batch_size] + # Bounded like every other provider call: this is blocking, and + # ``load_to_graph`` runs it inside the connect/refresh streams. response = batch_completion( model=Config.COMPLETION_MODEL, messages=batch_messages, temperature=0, max_tokens=50, + timeout=Config.LLM_TIMEOUT, + max_retries=Config.LLM_MAX_RETRIES, + num_retries=0, ) for offset, batch_response in enumerate(response): @@ -149,16 +155,16 @@ def generate_db_description( f"{tables_formatted}.\n\nDescription:" ) - response = completion( - model=Config.COMPLETION_MODEL, - messages=[ + # Via run_completion for the shared timeout, retry budget and duration + # logging. Blocking: async callers must offload it (see graph_loader). + return run_completion( + [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": prompt}, ], + label="db_description", temperature=temperature, max_tokens=max_tokens, n=1, stop=None, ) - description = response.choices[0].message["content"] - return description diff --git a/tests/test_embeddings_offloading.py b/tests/test_embeddings_offloading.py index 2ede3a34..ea943fed 100644 --- a/tests/test_embeddings_offloading.py +++ b/tests/test_embeddings_offloading.py @@ -100,6 +100,36 @@ def test_async_callers_do_not_embed_inline(): assert not offenders, "inline embedding call(s):\n" + "\n".join(offenders) +@pytest.mark.unit +def test_async_callers_do_not_call_llms_inline(): + """No bare provider call in modules whose coroutines back a stream. + + Every one of these has been a real incident-class bug: a synchronous + provider call inside an ``async def`` blocks the event loop, and a blocked + loop cannot write keepalives, so open streams are severed regardless of the + keepalive wrapper. + """ + import api.graph + import api.loaders.graph_loader as graph_loader + import api.memory.graphiti_tool as graphiti_tool + import api.routes.settings as settings_route + + # Bare provider entry points that must never be invoked on the loop. + calls = ("completion(", "batch_completion(", "embedding(") + offenders = [] + for module in (api.graph, graph_loader, graphiti_tool, settings_route): + source = inspect.getsource(module) + for lineno, line in enumerate(source.splitlines(), start=1): + stripped = line.strip() + if stripped.startswith("#") or "import" in stripped: + continue + if any(call in stripped for call in calls): + if "to_thread" not in stripped and "off_loop" not in stripped: + offenders.append(f"{module.__name__}:{lineno}: {stripped}") + + assert not offenders, "inline provider call(s):\n" + "\n".join(offenders) + + @pytest.mark.unit def test_embedding_calls_are_time_bounded(): """A hung provider must not pin a worker thread forever.""" diff --git a/tests/test_settings_route.py b/tests/test_settings_route.py index 1db6bc5d..9c23a654 100644 --- a/tests/test_settings_route.py +++ b/tests/test_settings_route.py @@ -2,6 +2,8 @@ from unittest.mock import patch, MagicMock +from api.config import Config + import pytest from api.routes.settings import validate_api_key, ValidateKeyRequest, _sanitize_for_log @@ -112,11 +114,16 @@ async def test_valid_key_returns_success(self, mock_completion, mock_request): body = response.body.decode() assert '"valid":true' in body + # The validation call is bounded like every other provider call, so a + # bad endpoint cannot hang the route (and with it every open stream). mock_completion.assert_called_once_with( model="openai/gpt-3.5-turbo", messages=[{"role": "user", "content": "test"}], max_tokens=1, api_key="sk-validkey123456", + timeout=Config.LLM_TIMEOUT, + max_retries=Config.LLM_MAX_RETRIES, + num_retries=0, ) @pytest.mark.asyncio @@ -163,11 +170,16 @@ async def test_gemini_vendor_accepted(self, mock_completion, mock_request): response = await validate_api_key.__wrapped__(mock_request, data) assert response.status_code == 200 + # The validation call is bounded like every other provider call, so a + # bad endpoint cannot hang the route (and with it every open stream). mock_completion.assert_called_once_with( model="gemini/gemini-pro", messages=[{"role": "user", "content": "test"}], max_tokens=1, api_key="AIzaSyTest123456", + timeout=Config.LLM_TIMEOUT, + max_retries=Config.LLM_MAX_RETRIES, + num_retries=0, ) @pytest.mark.asyncio From 1ba6cadeab0898ba38a69e39f47a50b65215aeb8 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Thu, 20 Aug 2026 15:55:24 +0300 Subject: [PATCH 15/25] fix: bound schema introspection, honour stricter timeout units, reject nan/inf MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fifth review from @Naseem77; all three findings were valid. **1. Cancellation could not release hung schema sessions.** Introspection had only a connect timeout, on the reasoning that a large schema may legitimately outlast the user-query ceiling. That was the wrong conclusion: the answer is a larger deadline, not none. A database that accepts the connection and then stalls held both its session and a worker thread, and since cancelling the awaiting task cannot stop that thread, repeated connect/refresh attempts could exhaust the executor every other offloaded call shares. Adds `DB_SCHEMA_TIMEOUT` (300s) — a server-side `statement_timeout` for PostgreSQL, socket read/write timeouts for MySQL, network and `STATEMENT_TIMEOUT_IN_SECONDS` for Snowflake — and `DB_SCHEMA_CONCURRENCY` (2), a semaphore in `api/loaders/introspection.py` capping how many introspections may hold workers at once. **2. Stricter PostgreSQL timeouts were being loosened.** The clamp only recognised lowercase bare digits, so a URL asking for `5s` was silently replaced with the 60s ceiling, and an uppercase `STATEMENT_TIMEOUT=` directive was not even stripped — it survived alongside ours, and GUC names are case-insensitive, so it could win. Values are now parsed case-insensitively with units (`us`/`ms`/`s`/`min`/`h`/`d`) and optional quotes, normalised to milliseconds, and the strictest positive value wins, capped at the configured ceiling. Sub-millisecond requests round up to 1ms rather than truncating to 0 and being discarded — which would have loosened them. **3. `nan` and `inf` passed validation.** Both slip past a `<= 0` test: nan compares False against everything and inf is a deadline that never expires. `math.isfinite` is now required. Tests: 13 new cases. The unit/case matrix, the duplicate and disabled shapes, the schema deadline on PostgreSQL and MySQL, and the concurrency cap. The last has teeth — widening the semaphore fails with "6 introspections ran concurrently, cap is 2". Two earlier clamp tests asserted duplicates collapse to the ceiling; they now assert the stricter, correct contract. 277 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- .env.example | 6 ++ api/config.py | 25 +++++++- api/loaders/introspection.py | 36 +++++++++++ api/loaders/mysql_loader.py | 12 +++- api/loaders/postgres_loader.py | 92 ++++++++++++++++++++-------- api/loaders/snowflake_loader.py | 12 +++- tests/test_db_execution_timeouts.py | 60 +++++++++++++----- tests/test_schema_load_offloading.py | 92 ++++++++++++++++++++++++++++ 8 files changed, 289 insertions(+), 46 deletions(-) create mode 100644 api/loaders/introspection.py diff --git a/.env.example b/.env.example index 0cd1bfd6..db55da11 100644 --- a/.env.example +++ b/.env.example @@ -94,6 +94,12 @@ FALKORDB_URL=redis://localhost:6379/0 # REQUIRED - change to your FalkorDB URL # cancelled from Python). Seconds. # DB_CONNECT_TIMEOUT=10 # DB_STATEMENT_TIMEOUT=60 +# +# Schema introspection gets a larger deadline (it is metadata work over a whole +# database) and a cap on how many may occupy worker threads at once, since that +# executor is shared with every other offloaded call. +# DB_SCHEMA_TIMEOUT=300 +# DB_SCHEMA_CONCURRENCY=2 # OpenAI - uses openai/gpt-4.1 and openai/text-embedding-ada-002 # OPENAI_API_KEY=your_openai_api_key diff --git a/api/config.py b/api/config.py index 63e645f9..d51ccb0d 100644 --- a/api/config.py +++ b/api/config.py @@ -6,6 +6,7 @@ import os import time import logging +import math import dataclasses from typing import Union @@ -127,6 +128,13 @@ def _positive_env(name: str, default: str, cast=int): raise ValueError( f"{name} must be a positive number (got {raw!r})" ) from exc + # ``nan`` and ``inf`` are floats that pass a ``<= 0`` test: nan compares + # False against everything, and inf is a deadline that never expires. + if not math.isfinite(value): + raise ValueError( + f"{name} must be a finite number (got {raw!r}); nan and inf are not " + "usable deadlines" + ) if value <= 0: raise ValueError( f"{name} must be greater than 0 (got {raw!r}); a zero or negative " @@ -136,7 +144,7 @@ def _positive_env(name: str, default: str, cast=int): @dataclasses.dataclass -class Config: +class Config: # pylint: disable=too-many-instance-attributes """ Configuration class for the text2sql module. """ @@ -224,6 +232,21 @@ class Config: # pylint: disable-next=invalid-name DB_STATEMENT_TIMEOUT: int = _positive_env("DB_STATEMENT_TIMEOUT", "60") + # Schema introspection gets its own, larger deadline: it is metadata work + # over a whole database, so the user-query ceiling is too tight, but it + # still needs a bound. Cancelling the awaiting task does not stop the + # driver call, so without this a stalled database holds both a session and + # a worker thread until it decides to answer. + # pylint: disable-next=invalid-name + DB_SCHEMA_TIMEOUT: int = _positive_env("DB_SCHEMA_TIMEOUT", "300") + + # How many schema introspections may occupy worker threads at once. The + # default executor is shared with every other offloaded call (LLM, + # embedding, user SQL), so unbounded schema work on a stalled database + # could starve all of it. + # pylint: disable-next=invalid-name + DB_SCHEMA_CONCURRENCY: int = _positive_env("DB_SCHEMA_CONCURRENCY", "2") + DB_MAX_DISTINCT: int = 100 # pylint: disable=invalid-name DB_UNIQUENESS_THRESHOLD: float = 0.5 # pylint: disable=invalid-name SHORT_MEMORY_LENGTH = 5 # Maximum number of questions to keep in short-term memory diff --git a/api/loaders/introspection.py b/api/loaders/introspection.py new file mode 100644 index 00000000..7e7f2bf7 --- /dev/null +++ b/api/loaders/introspection.py @@ -0,0 +1,36 @@ +"""Bounded execution for schema introspection. + +Introspection runs in a worker thread because the drivers are blocking, and +cancelling the awaiting task does not stop that thread: a stalled database +holds its session and its worker until it answers. Two bounds follow from +that — a deadline the server applies (see each loader's connect parameters), +and a cap on how many introspections may occupy the shared executor at once, +so a stalled database cannot starve every other offloaded call. +""" + +import asyncio +import logging + +from api.config import Config + +_SLOTS: asyncio.Semaphore | None = None + + +def _semaphore() -> asyncio.Semaphore: + """Create the semaphore lazily, on the loop that first needs it.""" + global _SLOTS # pylint: disable=global-statement + if _SLOTS is None: + _SLOTS = asyncio.Semaphore(Config.DB_SCHEMA_CONCURRENCY) + return _SLOTS + + +async def run_introspection(func, /, *args, **kwargs): + """Run *func* in a worker thread, holding one introspection slot.""" + slots = _semaphore() + if slots.locked(): + logging.info( + "schema introspection queued: %d concurrent slots in use", + Config.DB_SCHEMA_CONCURRENCY, + ) + async with slots: + return await asyncio.to_thread(func, *args, **kwargs) diff --git a/api/loaders/mysql_loader.py b/api/loaders/mysql_loader.py index e81d82c6..fdeaeb19 100644 --- a/api/loaders/mysql_loader.py +++ b/api/loaders/mysql_loader.py @@ -1,6 +1,5 @@ """MySQL loader for loading database schemas into FalkorDB graphs.""" -import asyncio import datetime import decimal import logging @@ -15,6 +14,7 @@ from api.config import Config from api.loaders.base_loader import BaseLoader from api.loaders.graph_loader import load_to_graph +from api.loaders.introspection import run_introspection class MySQLQueryError(Exception): @@ -170,8 +170,14 @@ def _introspect_schema(conn_params: Dict[str, Any], db_name: str): conn = None cursor = None try: + # Socket deadlines for the introspection itself, larger than the + # user-query ceiling but still bounded: cancelling the awaiting + # task cannot stop this thread. conn = pymysql.connect( - connect_timeout=Config.DB_CONNECT_TIMEOUT, **conn_params + connect_timeout=Config.DB_CONNECT_TIMEOUT, + read_timeout=Config.DB_SCHEMA_TIMEOUT, + write_timeout=Config.DB_SCHEMA_TIMEOUT, + **conn_params, ) cursor = conn.cursor(DictCursor) entities = MySQLLoader.extract_tables_info(cursor, db_name) @@ -206,7 +212,7 @@ async def load( # pylint: disable=arguments-differ db_name = conn_params['database'] yield True, "Extracting table information..." - entities, relationships = await asyncio.to_thread( + entities, relationships = await run_introspection( MySQLLoader._introspect_schema, conn_params, db_name ) diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index 1239bdc9..f5338508 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -1,6 +1,5 @@ """PostgreSQL loader for loading database schemas into FalkorDB graphs.""" -import asyncio import re import datetime import decimal @@ -15,6 +14,14 @@ from api.config import Config from api.loaders.base_loader import BaseLoader # pylint: disable=import-error from api.loaders.graph_loader import load_to_graph # pylint: disable=import-error +from api.loaders.introspection import run_introspection + +# A real ``-c statement_timeout=`` directive, case-insensitive (GUC +# names are), capturing an optionally quoted value that may carry a unit. +_STATEMENT_TIMEOUT_RE = re.compile( + r"(?i)(?:^|\s)-c\s*statement_timeout\s*=\s*" + r"('[^']*'|\"[^\"]*\"|\S*)" +) logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") @@ -141,6 +148,37 @@ def parse_schema_from_url(connection_url: str) -> str: except Exception: # pylint: disable=broad-exception-caught return 'public' + @staticmethod + def _statement_timeout_ms(raw: str): + """Normalise a ``statement_timeout`` value to milliseconds. + + PostgreSQL accepts a bare number (milliseconds) or a number with a unit + (``us``, ``ms``, ``s``, ``min``, ``h``, ``d``), optionally quoted. A + URL asking for ``5s`` is stricter than a 60s ceiling and should be + honoured, so the value has to be understood rather than pattern-matched + as digits. Returns ``None`` when the value is not something we can + compare, in which case the configured ceiling is used instead. + """ + value = raw.strip().strip('"\'') + match = re.fullmatch( + r"(?i)\s*(\d+(?:\.\d+)?)\s*(us|ms|s|min|h|d)?\s*", value + ) + if not match: + return None + amount = float(match.group(1)) + unit = (match.group(2) or "ms").lower() + factors = { + "us": 0.001, "ms": 1, "s": 1000, + "min": 60_000, "h": 3_600_000, "d": 86_400_000, + } + milliseconds = amount * factors[unit] + if milliseconds <= 0: + return None + # Round up so a sub-millisecond request (e.g. ``500us``) is honoured as + # the strictest representable bound rather than truncated to 0 and + # discarded, which would silently loosen it to the ceiling. + return max(1, int(milliseconds)) + @staticmethod def _introspect_schema(connection_url: str, schema: str): """Connect, introspect and close — all inside one worker thread. @@ -160,8 +198,14 @@ def _introspect_schema(connection_url: str, schema: str): conn = None cursor = None try: + # A server-side deadline for the introspection itself, larger + # than the user-query ceiling but still bounded: cancelling the + # awaiting task cannot stop this thread, so a stalled database + # would otherwise hold the session and the worker indefinitely. conn = psycopg2.connect( - connection_url, connect_timeout=Config.DB_CONNECT_TIMEOUT + connection_url, + connect_timeout=Config.DB_CONNECT_TIMEOUT, + options=f"-c statement_timeout={Config.DB_SCHEMA_TIMEOUT * 1000}", ) cursor = conn.cursor() @@ -209,7 +253,7 @@ async def load( # pylint: disable=arguments-differ db_name = db_name.split('?')[0] yield True, "Extracting table information..." - entities, relationships = await asyncio.to_thread( + entities, relationships = await run_introspection( PostgresLoader._introspect_schema, connection_url, schema ) @@ -567,32 +611,28 @@ def _execution_connect_kwargs(db_url: str) -> Dict[str, Any]: url_options = url_params.get("options", [""])[0] kwargs: Dict[str, Any] = {} - # The configured values are maximums, not defaults: a URL may tighten - # them but must not loosen them or switch them off. In libpq a timeout - # of 0 means "no limit", which would let one query hold an - # uncancellable worker thread indefinitely. + # The configured value is a maximum, not a default: a URL may tighten + # it but must not loosen it or switch it off. In libpq a timeout of 0 + # means "no limit", which would let one query hold an uncancellable + # worker thread indefinitely. # - # Every accepted directive is removed and exactly one normalised bound - # appended. Leaving any in place is not safe: libpq applies the last - # occurrence, so `statement_timeout=1000 ... statement_timeout=0` would - # end up unbounded, and a unit-bearing value like `2min` is not - # comparable to our millisecond ceiling. - # - # The pattern matches a real directive, not the bare word: a substring - # test would also hit `-c application_name=statement_timeout_probe`. + # Every directive is removed and exactly one canonical bound appended. + # Leaving any in place is unsafe: libpq applies the last occurrence, so + # `statement_timeout=1000 ... statement_timeout=0` would end up + # unbounded. GUC names are case-insensitive, so an uppercase directive + # left behind would win over ours. timeout_ms = Config.DB_STATEMENT_TIMEOUT * 1000 - directive = re.compile(r"(?:^|\s)-c\s*statement_timeout\s*=\s*(\S*)") - found = directive.findall(url_options) - stripped = directive.sub(" ", url_options).strip() - - # Honour a URL value only when it is unambiguous: a single directive, - # plain milliseconds, and no looser than the configured ceiling. - effective_ms = timeout_ms - if len(found) == 1 and found[0].isdigit(): - requested = int(found[0]) - if 0 < requested <= timeout_ms: - effective_ms = requested + requested = [ + ms for ms in ( + PostgresLoader._statement_timeout_ms(raw) + for raw in _STATEMENT_TIMEOUT_RE.findall(url_options) + ) + if ms is not None and ms > 0 + ] + # Strictest wins, and never looser than the configured ceiling. + effective_ms = min([timeout_ms, *requested]) + stripped = _STATEMENT_TIMEOUT_RE.sub(" ", url_options).strip() options = f"{stripped} -c statement_timeout={effective_ms}".strip() if options: kwargs["options"] = options diff --git a/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index 936e04df..13daa92b 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -1,6 +1,5 @@ """Snowflake loader for loading database schemas into FalkorDB graphs.""" -import asyncio import base64 import datetime import decimal @@ -19,6 +18,7 @@ from api.config import Config from api.loaders.base_loader import BaseLoader from api.loaders.graph_loader import load_to_graph +from api.loaders.introspection import run_introspection class SnowflakeQueryError(Exception): @@ -257,6 +257,14 @@ def _introspect_schema( conn = None cursor = None try: + # Bound the introspection itself, not just the login: cancelling + # the awaiting task cannot stop this thread. + conn_params = dict(conn_params) + conn_params["network_timeout"] = Config.DB_SCHEMA_TIMEOUT + conn_params["session_parameters"] = { + **(conn_params.get("session_parameters") or {}), + "STATEMENT_TIMEOUT_IN_SECONDS": Config.DB_SCHEMA_TIMEOUT, + } conn = snowflake.connector.connect(**conn_params) cursor = conn.cursor(DictCursor) entities = SnowflakeLoader.extract_tables_info( @@ -295,7 +303,7 @@ async def load(prefix: str, connection_url: str) -> AsyncGenerator[ schema_name = conn_params['schema'].upper() yield True, "Extracting table information..." - entities, relationships = await asyncio.to_thread( + entities, relationships = await run_introspection( SnowflakeLoader._introspect_schema, conn_params, db_name, schema_name, ) diff --git a/tests/test_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py index 3445fa53..596d3b92 100644 --- a/tests/test_db_execution_timeouts.py +++ b/tests/test_db_execution_timeouts.py @@ -133,24 +133,56 @@ def test_snowflake_overrides_parser_timeout_defaults(mock_connect): @pytest.mark.unit -@pytest.mark.parametrize("url_options,reason", [ - ("-c%20statement_timeout%3D1000%20-c%20statement_timeout%3D0", "duplicate, last disables"), - ("-c%20statement_timeout%3D0%20-c%20statement_timeout%3D1000", "duplicate, first disables"), - ("-c%20statement_timeout%3D2min", "unit-bearing value"), - ("-c%20statement_timeout%3D%20", "empty value"), - ("-c%20statement_timeout%3D-5", "negative value"), +@pytest.mark.parametrize("url_options,expected_ms,reason", [ + # libpq applies the last directive, so none may survive; the strictest + # positive value wins and a disabling 0 is ignored entirely. + ("-c%20statement_timeout%3D1000%20-c%20statement_timeout%3D0", 1000, + "duplicate, last disables"), + ("-c%20statement_timeout%3D0%20-c%20statement_timeout%3D1000", 1000, + "duplicate, first disables"), + ("-c%20STATEMENT_TIMEOUT%3D0%20-c%20statement_timeout%3D3000", 3000, + "mixed case duplicate"), + ("-c%20statement_timeout%3D2min", None, "looser unit value"), + ("-c%20statement_timeout%3D0", None, "disabled"), + ("-c%20statement_timeout%3D%20", None, "empty value"), + ("-c%20statement_timeout%3D-5", None, "negative value"), + ("-c%20statement_timeout%3D0s", None, "zero with a unit"), ]) -def test_postgres_clamp_is_not_bypassable(url_options, reason): - """libpq applies the last directive, so none may survive. +def test_postgres_clamp_is_not_bypassable(url_options, expected_ms, reason): + """Exactly one directive survives, never looser than the ceiling. - Leaving an accepted directive in place lets ``statement_timeout=1000 ... - statement_timeout=0`` end up unbounded, and a unit-bearing value like - ``2min`` is not comparable to the millisecond ceiling. + ``expected_ms=None`` means the configured ceiling applies. """ + ceiling = Config.DB_STATEMENT_TIMEOUT * 1000 kwargs = PostgresLoader._execution_connect_kwargs(f"{PG_URL}?options={url_options}") - expected = f"-c statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}" - assert kwargs["options"] == expected, reason - assert kwargs["options"].count("statement_timeout") == 1, reason + options = kwargs["options"] + + assert options.lower().count("statement_timeout") == 1, reason + assert options == f"-c statement_timeout={expected_ms or ceiling}", reason + + +@pytest.mark.unit +@pytest.mark.parametrize("value,expected_ms", [ + ("5s", 5_000), # stricter than the 60s ceiling, so honoured + ("5000", 5_000), # bare numbers are milliseconds + ("'5s'", 5_000), # quoted + ("500us", 1), # sub-millisecond rounds up rather than truncating + ("1min", 60_000), # equal to the ceiling + ("2min", None), # looser, so clamped +]) +def test_postgres_honours_stricter_units_and_case(value, expected_ms): + """PostgreSQL accepts units, quotes and any case; all must be understood. + + Treating only lowercase bare digits as valid silently loosened a URL asking + for ``5s`` to the configured 60s. + """ + ceiling = Config.DB_STATEMENT_TIMEOUT * 1000 + quoted = value.replace("'", "%27").replace(" ", "%20") + for name in ("statement_timeout", "STATEMENT_TIMEOUT", "Statement_Timeout"): + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options=-c%20{name}%3D{quoted}" + ) + assert kwargs["options"] == f"-c statement_timeout={expected_ms or ceiling}", name @pytest.mark.unit diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py index 8bb0e9c3..1787f1fd 100644 --- a/tests/test_schema_load_offloading.py +++ b/tests/test_schema_load_offloading.py @@ -7,11 +7,13 @@ """ import asyncio +import threading import time from unittest.mock import MagicMock, patch import pytest +from api.config import Config from api.core.pipeline import MySQLLoader, PostgresLoader STALL = 0.3 @@ -183,3 +185,93 @@ def boom(*_args, **_kwargs): # ...and the connection is closed regardless. assert conn.close.called, "connection leaked when introspection failed" + + +@pytest.mark.unit +@patch("api.loaders.postgres_loader.load_to_graph") +@patch("api.loaders.postgres_loader.PostgresLoader.extract_relationships", _slow) +@patch("api.loaders.postgres_loader.PostgresLoader.extract_tables_info", _slow) +@patch("api.loaders.postgres_loader.psycopg2.connect") +async def test_postgres_schema_introspection_is_time_bounded( + mock_connect, mock_load_to_graph +): + """Introspection carries its own, larger server-side deadline. + + A connect timeout alone is not enough: cancelling the awaiting task cannot + stop the driver call, so a database that accepts the connection and then + stalls would hold the session and the worker until it answered. + """ + async def noop(*_args, **_kwargs): + return None + + mock_load_to_graph.side_effect = noop + mock_connect.return_value = MagicMock() + + async for _ in PostgresLoader.load("pfx", "postgresql://u:p@h:5432/db"): + pass + + kwargs = mock_connect.call_args.kwargs + assert kwargs["connect_timeout"] == Config.DB_CONNECT_TIMEOUT + expected = f"-c statement_timeout={Config.DB_SCHEMA_TIMEOUT * 1000}" + assert kwargs["options"] == expected + # Deliberately larger than the user-query ceiling: metadata work over a + # whole database legitimately outlasts a single query. + assert Config.DB_SCHEMA_TIMEOUT > Config.DB_STATEMENT_TIMEOUT + + +@pytest.mark.unit +@patch("api.loaders.mysql_loader.pymysql.connect") +async def test_mysql_schema_introspection_is_time_bounded(mock_connect): + cursor = MagicMock() + cursor.description = None + mock_connect.return_value.cursor.return_value = cursor + + with patch.object(MySQLLoader, "extract_tables_info", lambda *_a: {}), \ + patch.object(MySQLLoader, "extract_relationships", lambda *_a: {}), \ + patch("api.loaders.mysql_loader.load_to_graph") as load_to_graph: + async def noop(*_args, **_kwargs): + return None + + load_to_graph.side_effect = noop + async for _ in MySQLLoader.load("pfx", "mysql://u:p@h:3306/db"): + pass + + kwargs = mock_connect.call_args.kwargs + assert kwargs["connect_timeout"] == Config.DB_CONNECT_TIMEOUT + assert kwargs["read_timeout"] == Config.DB_SCHEMA_TIMEOUT + assert kwargs["write_timeout"] == Config.DB_SCHEMA_TIMEOUT + + +@pytest.mark.unit +async def test_schema_introspection_concurrency_is_bounded(monkeypatch): + """Only DB_SCHEMA_CONCURRENCY introspections may hold workers at once. + + The executor is shared with every other offloaded call, so unbounded schema + work against a stalled database could starve LLM, embedding and user-SQL + calls alike. + """ + import api.loaders.introspection as introspection + + monkeypatch.setattr(introspection, "_SLOTS", None) + monkeypatch.setattr(Config, "DB_SCHEMA_CONCURRENCY", 2, raising=False) + + live = 0 + peak = 0 + lock = threading.Lock() + + def blocking_work(): + nonlocal live, peak + with lock: + live += 1 + peak = max(peak, live) + time.sleep(STALL) + with lock: + live -= 1 + return "done" + + results = await asyncio.gather( + *(introspection.run_introspection(blocking_work) for _ in range(6)) + ) + + assert results == ["done"] * 6 + assert peak <= 2, f"{peak} introspections ran concurrently, cap is 2" From a14e3ccbc420eeef8af4ab19d3d46d6770514eb6 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 09:48:05 +0300 Subject: [PATCH 16/25] fix(loaders): recognise PostgreSQL's long-option timeout directive form MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up on @Naseem77's review item about "long-option forms", which I had missed: PostgreSQL accepts `--name=value` in an options string as well as `-c name=value`, and converts hyphens in the long form to underscores. The pattern only matched the short form, so `--statement_timeout=...` and `--statement-timeout=...` survived in the string. The disabling case was not exploitable in practice — our directive is appended last and PostgreSQL applies the last assignment — but the stricter case was wrong in the direction that matters: `--statement_timeout=5s` was left in place and then overridden by the configured 60s ceiling, silently loosening the request. Both forms are now matched, case-insensitively, with either separator. Tests parametrise all five spellings (`-c name=`, `-cname=`, `--name=`, `--name-with-hyphens=`, uppercase) and assert both directions: a disabling value never survives, and a stricter value is honoured. Reverting to the short-form-only pattern fails three of them. 280 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/loaders/postgres_loader.py | 7 ++++++- tests/test_db_execution_timeouts.py | 30 +++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index f5338508..2052ec7d 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -19,7 +19,12 @@ # A real ``-c statement_timeout=`` directive, case-insensitive (GUC # names are), capturing an optionally quoted value that may carry a unit. _STATEMENT_TIMEOUT_RE = re.compile( - r"(?i)(?:^|\s)-c\s*statement_timeout\s*=\s*" + # PostgreSQL accepts both directive forms in an options string: + # ``-c name=value`` and ``--name=value``, the latter also with hyphens in + # place of underscores. GUC names are case-insensitive. Missing a form + # leaves it in the string, where it either overrides our bound or, if ours + # wins, silently loosens a stricter request. + r"(?i)(?:^|\s)(?:-c\s*|--)statement[-_]timeout\s*=\s*" r"('[^']*'|\"[^\"]*\"|\S*)" ) diff --git a/tests/test_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py index 596d3b92..879fd2dc 100644 --- a/tests/test_db_execution_timeouts.py +++ b/tests/test_db_execution_timeouts.py @@ -161,6 +161,36 @@ def test_postgres_clamp_is_not_bypassable(url_options, expected_ms, reason): assert options == f"-c statement_timeout={expected_ms or ceiling}", reason +@pytest.mark.unit +@pytest.mark.parametrize("directive", [ + "-c%20statement_timeout", # short form + "-cstatement_timeout", # short form, no space + "--statement_timeout", # long form + "--statement-timeout", # long form, hyphenated + "--STATEMENT-TIMEOUT", # long form, hyphenated, uppercase +]) +def test_postgres_recognises_every_directive_form(directive): + """PostgreSQL accepts ``-c name=value`` and ``--name=value``. + + A form we do not recognise stays in the options string, where it either + overrides our bound or — since ours is appended last — silently loosens a + stricter request. + """ + ceiling = Config.DB_STATEMENT_TIMEOUT * 1000 + + # A disabling value must never survive. + disabled = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options={directive}%3D0" + )["options"] + assert disabled == f"-c statement_timeout={ceiling}", directive + + # A stricter value must be honoured, whichever form it arrives in. + stricter = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options={directive}%3D5s" + )["options"] + assert stricter == "-c statement_timeout=5000", directive + + @pytest.mark.unit @pytest.mark.parametrize("value,expected_ms", [ ("5s", 5_000), # stricter than the 60s ceiling, so honoured From 938f70cd929c6ab122da8142a918a0dd7fad6996 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 10:19:37 +0300 Subject: [PATCH 17/25] fix: propagate producer cancellation, hold introspection slots, unbreak Snowflake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three of the four open review threads turned out to be real, and two of them were defects in code this PR added. **1. `with_keepalive` hung instead of ending (Copilot).** `CancelledError` is a `BaseException`, so the pump never relays it as a queue item; the consumer then sat on an empty queue emitting keepalives forever. Reproduced: 39 keepalives and still going after 2s. The consumer now observes the pump's terminal state, drains anything already enqueued, and re-raises via `pump.result()`, so cancellation reaches the caller. This was a consequence of the earlier, correct switch from `except BaseException` to `except Exception`. **2. Introspection slots were released on cancellation (Copilot).** Cancelling the awaiting task cannot stop a running worker, so the slot was handed to new work while the old thread still held a worker and a database session — exactly the disconnect-driven load the cap exists to bound. Reproduced: peak 4 workers against a cap of 2. The slot is now released by the worker's done-callback, and the worker is shielded so cancellation still propagates while it keeps its slot. **3. Snowflake could not be connected or refreshed at all (Copilot).** `schema_loader` calls `loader.load(user_id, url, db=db)` and `_emit_schema_refresh` calls `refresh_graph_schema(..., db=db)`, but the Snowflake versions took no `db` parameter, so both raised `TypeError`. Refresh also reached for the `api.extensions` singleton instead of the caller's handle. Both now match the PostgreSQL and MySQL signatures and use `resolve_db(db)`. Pre-existing on staging and unrelated to the incident, so strictly out of scope for this PR — fixed here because it is a two-line crash on a supported database, and `tests/test_loader_contract.py` now pins the signature contract across all three loaders so it cannot silently drift again. The fourth open thread (CodeRabbit, total-call deadline in `api/config.py`) remains deliberately open: see the earlier reply for why a hard wall-clock deadline is not enforceable for calls running in `asyncio.to_thread`. Tests: regression cases for the hang and for the slot, plus 9 contract cases. Both new tests were checked against the pre-fix behaviour. 293 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/loaders/introspection.py | 31 ++++++++++++++++-- api/loaders/snowflake_loader.py | 14 ++++---- api/routes/streaming.py | 19 +++++++++++ tests/test_loader_contract.py | 49 ++++++++++++++++++++++++++++ tests/test_schema_load_offloading.py | 48 +++++++++++++++++++++++++++ tests/test_stream_keepalive.py | 26 +++++++++++++++ 6 files changed, 178 insertions(+), 9 deletions(-) create mode 100644 tests/test_loader_contract.py diff --git a/api/loaders/introspection.py b/api/loaders/introspection.py index 7e7f2bf7..2a15e33e 100644 --- a/api/loaders/introspection.py +++ b/api/loaders/introspection.py @@ -25,12 +25,37 @@ def _semaphore() -> asyncio.Semaphore: async def run_introspection(func, /, *args, **kwargs): - """Run *func* in a worker thread, holding one introspection slot.""" + """Run *func* in a worker thread, holding one introspection slot. + + The slot is released when the *thread* finishes, not when the awaiting task + ends. Cancelling the awaiting task cannot stop a running worker, so + releasing on cancellation would hand the slot to new work while the old + thread still holds a worker and a database session — letting the cap be + exceeded by exactly the disconnect-driven load it exists to bound. + + Cancellation still reaches the caller: the worker is shielded so it keeps + running (and keeps its slot) while the ``CancelledError`` propagates. + """ slots = _semaphore() if slots.locked(): logging.info( "schema introspection queued: %d concurrent slots in use", Config.DB_SCHEMA_CONCURRENCY, ) - async with slots: - return await asyncio.to_thread(func, *args, **kwargs) + + await slots.acquire() + try: + worker = asyncio.ensure_future(asyncio.to_thread(func, *args, **kwargs)) + except BaseException: + slots.release() + raise + + def _release(finished: asyncio.Future) -> None: + slots.release() + if not finished.cancelled(): + # Mark any failure retrieved: when the caller was cancelled nobody + # is left to await this future. + finished.exception() + + worker.add_done_callback(_release) + return await asyncio.shield(worker) diff --git a/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index 13daa92b..4adc0d65 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -281,7 +281,7 @@ def _introspect_schema( conn.close() @staticmethod - async def load(prefix: str, connection_url: str) -> AsyncGenerator[ + async def load(prefix: str, connection_url: str, db=None) -> AsyncGenerator[ tuple[bool, str], None ]: """ @@ -311,7 +311,7 @@ async def load(prefix: str, connection_url: str) -> AsyncGenerator[ # Load data into graph yield True, "Loading data into graph..." await load_to_graph(f"{prefix}_{db_name}", entities, relationships, - db_name=db_name, db_url=connection_url) + db_name=db_name, db_url=connection_url, db=db) yield True, (f"Snowflake schema loaded successfully. " f"Found {len(entities)} tables.") @@ -606,7 +606,9 @@ def is_schema_modifying_query(sql_query: str) -> Tuple[bool, str]: return False, "" @staticmethod - async def refresh_graph_schema(graph_id: str, db_url: str) -> Tuple[bool, str]: + async def refresh_graph_schema( + graph_id: str, db_url: str, db=None + ) -> Tuple[bool, str]: """ Refresh the graph schema by clearing existing data and reloading from the database. @@ -621,11 +623,11 @@ async def refresh_graph_schema(graph_id: str, db_url: str) -> Tuple[bool, str]: logging.info("Schema modification detected. Refreshing graph schema.") # Import here to avoid circular imports - from api.extensions import db # pylint: disable=import-error,import-outside-toplevel + from api.core.db_resolver import resolve_db # pylint: disable=import-outside-toplevel # Clear existing graph data # Drop current graph before reloading - graph = db.select_graph(graph_id) + graph = resolve_db(db).select_graph(graph_id) await graph.delete() # Extract prefix from graph_id (remove database name part) @@ -640,7 +642,7 @@ async def refresh_graph_schema(graph_id: str, db_url: str) -> Tuple[bool, str]: # Reuse the existing load method to reload the schema success = False message = "" - async for progress_tuple in SnowflakeLoader.load(prefix, db_url): + async for progress_tuple in SnowflakeLoader.load(prefix, db_url, db=db): success, message = progress_tuple if success: diff --git a/api/routes/streaming.py b/api/routes/streaming.py index dfdfe2af..cada6d37 100644 --- a/api/routes/streaming.py +++ b/api/routes/streaming.py @@ -66,6 +66,25 @@ async def _pump(): try: item = await asyncio.wait_for(queue.get(), timeout=interval) except asyncio.TimeoutError: + if pump.done(): + # The producer ended without a terminal item, so it was + # cancelled — every other exit enqueues one. Emitting + # keepalives from here would never terminate the response. + # Drain anything it managed to enqueue first: the terminal + # item can land between the timeout firing and this check. + while not queue.empty(): + queued = queue.get_nowait() + if queued is finished: + return + if isinstance(queued, Exception): + # ``from None``: the timeout is how we noticed, not + # the cause. Chaining it would misreport the error. + raise queued from None + yield queued + # Re-raises the producer's CancelledError, so cancellation + # reaches the consumer instead of stalling it. + pump.result() + return yield MESSAGE_DELIMITER continue if item is finished: diff --git a/tests/test_loader_contract.py b/tests/test_loader_contract.py new file mode 100644 index 00000000..81231667 --- /dev/null +++ b/tests/test_loader_contract.py @@ -0,0 +1,49 @@ +"""Every loader must accept the explicit graph handle. + +``schema_loader`` calls ``loader.load(user_id, url, db=db)`` and +``_emit_schema_refresh`` calls ``loader.refresh_graph_schema(..., db=db)``, so a +loader missing the parameter raises ``TypeError`` at runtime rather than at +import — Snowflake did, which broke connecting and refreshing a Snowflake +database outright. A signature check catches that without needing a live +warehouse. +""" + +import inspect + +import pytest + +import api.core # noqa: F401 pylint: disable=unused-import +from api.loaders.mysql_loader import MySQLLoader +from api.loaders.postgres_loader import PostgresLoader +from api.loaders.snowflake_loader import SnowflakeLoader + +LOADERS = [PostgresLoader, MySQLLoader, SnowflakeLoader] + + +@pytest.mark.unit +@pytest.mark.parametrize("loader", LOADERS, ids=lambda l: l.__name__) +@pytest.mark.parametrize("method", ["load", "refresh_graph_schema"]) +def test_loader_accepts_explicit_db_handle(loader, method): + params = inspect.signature(getattr(loader, method)).parameters + assert "db" in params, ( + f"{loader.__name__}.{method} must accept db= — callers pass it, so a " + "missing parameter is a runtime TypeError" + ) + assert params["db"].default is None, ( + f"{loader.__name__}.{method} db= must be optional" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("loader", LOADERS, ids=lambda l: l.__name__) +def test_loader_refresh_resolves_the_handle(loader): + """Refresh must resolve the passed handle, not reach for the singleton. + + Using the module-level ``api.extensions.db`` ignores the caller's handle, + which is what the parameter exists to provide. + """ + source = inspect.getsource(getattr(loader, "refresh_graph_schema")) + assert "resolve_db(" in source, f"{loader.__name__} should use resolve_db(db)" + assert "from api.extensions import db" not in source, ( + f"{loader.__name__} refresh ignores the caller's handle" + ) diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py index 1787f1fd..b4b0fd42 100644 --- a/tests/test_schema_load_offloading.py +++ b/tests/test_schema_load_offloading.py @@ -275,3 +275,51 @@ def blocking_work(): assert results == ["done"] * 6 assert peak <= 2, f"{peak} introspections ran concurrently, cap is 2" + + +@pytest.mark.unit +async def test_cancelled_introspection_keeps_its_slot(monkeypatch): + """A cancelled introspection must not hand its slot to new work. + + Cancelling the awaiting task cannot stop the worker, so releasing the slot + on cancellation lets the cap be exceeded by exactly the disconnect-driven + load it exists to bound. + """ + import api.loaders.introspection as introspection + + monkeypatch.setattr(introspection, "_SLOTS", None) + monkeypatch.setattr(Config, "DB_SCHEMA_CONCURRENCY", 2, raising=False) + + live = 0 + peak = 0 + lock = threading.Lock() + + def blocking_work(): + nonlocal live, peak + with lock: + live += 1 + peak = max(peak, live) + time.sleep(STALL * 2) + with lock: + live -= 1 + + # Fill the cap, then cancel both awaiting tasks while the threads run on. + first = [ + asyncio.ensure_future(introspection.run_introspection(blocking_work)) + for _ in range(2) + ] + await asyncio.sleep(STALL / 2) + for task in first: + task.cancel() + await asyncio.gather(*first, return_exceptions=True) + + # Slots must still be held by the running threads. + second = [ + asyncio.ensure_future(introspection.run_introspection(blocking_work)) + for _ in range(4) + ] + await asyncio.sleep(STALL) + assert peak <= 2, f"{peak} workers ran concurrently while the cap was 2" + + await asyncio.gather(*second, return_exceptions=True) + diff --git a/tests/test_stream_keepalive.py b/tests/test_stream_keepalive.py index a41c5469..e0466150 100644 --- a/tests/test_stream_keepalive.py +++ b/tests/test_stream_keepalive.py @@ -172,3 +172,29 @@ async def source(): await agen.aclose() await asyncio.wait_for(cancelled.wait(), timeout=2) + + +@pytest.mark.unit +async def test_producer_cancellation_reaches_the_consumer(): + """A cancelling inner stream must end the response, not stall it. + + ``CancelledError`` is a ``BaseException``, so the pump does not relay it as + a queue item. Without observing the pump's terminal state the consumer sat + on an empty queue emitting keepalives forever — an endless response. + """ + async def source(): + yield "first" + raise asyncio.CancelledError() + + chunks = [] + + async def consume(): + async for chunk in with_keepalive(source(), interval=0.02): + chunks.append(chunk) + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(consume(), timeout=3) + + assert chunks[0] == "first" + # A handful of keepalives during the gap is fine; an unbounded stream is not. + assert len(chunks) < 50, "consumer kept emitting keepalives after the producer died" From 93bd93eee1c7d083cad6ad7234448fa920bed611 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 10:27:10 +0300 Subject: [PATCH 18/25] test: drop the side-effect import from the loader contract test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Import the loaders through api.core.pipeline, which initialises the package in the right order, instead of a bare `import api.core` for its side effect. The shim was load-bearing — removing it outright, as the scanner suggested, breaks collection with the pre-existing circular import — so this routes around it rather than deleting it. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_loader_contract.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_loader_contract.py b/tests/test_loader_contract.py index 81231667..ecdce4f0 100644 --- a/tests/test_loader_contract.py +++ b/tests/test_loader_contract.py @@ -12,9 +12,11 @@ import pytest -import api.core # noqa: F401 pylint: disable=unused-import -from api.loaders.mysql_loader import MySQLLoader -from api.loaders.postgres_loader import PostgresLoader +# Imported via api.core.pipeline: importing a loader module first hits a +# circular import (pipeline imports the loaders, the loaders import api.core). +# Going through pipeline initialises the package in the right order, which also +# makes the snowflake import below work. +from api.core.pipeline import MySQLLoader, PostgresLoader from api.loaders.snowflake_loader import SnowflakeLoader LOADERS = [PostgresLoader, MySQLLoader, SnowflakeLoader] From a0a90c920c7d875299bcb5482000391b138bfd54 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 10:38:33 +0300 Subject: [PATCH 19/25] fix(agents): make LLM_TIMEOUT a total-call budget, not a per-attempt one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the one review item I had been declining (CodeRabbit, api/config.py). My argument was that a hard wall-clock deadline is not enforceable for calls running in `asyncio.to_thread`, since Python cannot cancel a thread blocked in a socket read. That is still true, but it was the wrong conclusion: the deadline does not have to be enforced *outside* the call, it can be divided *across* the attempts inside it. `Config.llm_call_bounds()` now derives the per-attempt timeout from the total budget: `LLM_TIMEOUT / (LLM_MAX_RETRIES + 1)`, with litellm's own retry loop still disabled so the two cannot compound. One place defines the ceiling and both the completion and embedding paths use it, so they cannot drift apart. Measured against a local server that accepts the request and never replies: total=8s retries=1 -> per-attempt 4.0s -> aborted after 4.19s total=8s retries=0 -> per-attempt 8.0s -> aborted after 8.19s total=8s retries=3 -> per-attempt 2.0s -> aborted after 2.19s All within the budget. Before the retry budget was pinned at all, a 3s timeout took 10.81s to fail. Also addresses the second half of that review comment — `**kwargs` silently bypassing the bounds. An explicit call-site override is still allowed, since passing `timeout=` is deliberate, but it is now logged (`bound overrides in effect: {...}`) so it cannot weaken the ceiling unnoticed. No caller currently overrides. `LLM_TIMEOUT` changes meaning from per-attempt to total, so `.env.example` is updated and two tests that pinned the old semantics now assert the new contract. 298 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- .env.example | 8 +++-- api/agents/utils.py | 28 +++++++++------ api/config.py | 24 ++++++++++--- tests/test_embeddings_offloading.py | 9 +++-- tests/test_timeout_validation.py | 55 +++++++++++++++++++++++++++++ 5 files changed, 103 insertions(+), 21 deletions(-) diff --git a/.env.example b/.env.example index db55da11..8f5f1826 100644 --- a/.env.example +++ b/.env.example @@ -75,9 +75,11 @@ FALKORDB_URL=redis://localhost:6379/0 # REQUIRED - change to your FalkorDB URL # COMPLETION_MODEL=openai/gpt-4.1 # EMBEDDING_MODEL=openai/text-embedding-ada-002 -# Wall-clock ceiling for a single agent LLM call, in seconds (default 90). -# Passed to litellm, which aborts the HTTP request — a hung provider then -# surfaces as an error instead of stalling the response stream. +# Wall-clock budget for one LLM call end to end, in seconds (default 90). +# This is the total, not per attempt: the per-attempt timeout handed to the +# provider is this divided by LLM_MAX_RETRIES + 1, so retries cannot push the +# real ceiling past it. A hung provider surfaces as an error within the budget +# instead of stalling the response stream. # LLM_TIMEOUT=90 # # Calls slower than this are logged at WARNING (default 20). diff --git a/api/agents/utils.py b/api/agents/utils.py index aedc0099..78da2951 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -14,24 +14,30 @@ def run_completion(messages: List[Dict[str, str]], custom_model: str | None = No **kwargs) -> str: """Run an LLM completion with optional custom model/key overrides. - Applies ``Config.LLM_TIMEOUT`` per attempt and a pinned retry budget - unless the caller overrides them, and logs the call duration. Both exist because the 2026-07-29 - demo failure was an LLM call that stalled with no timeout and left no - trace of how long it ran. ``label`` names the caller in those log lines - and is not forwarded to the provider. + Bounds the call with ``Config.llm_call_bounds()``: ``LLM_TIMEOUT`` is the + budget for the whole call, divided across attempts, so retries cannot push + the real ceiling past it. Duration is logged. Both exist because the + 2026-07-29 demo failure was an LLM call that stalled with no timeout and + left no trace of how long it ran. ``label`` names the caller in those log + lines and is not forwarded to the provider. + + A caller may still override the bounds explicitly; doing so is logged so it + cannot silently weaken the ceiling. Returns the content string from the first choice. """ + bounds = Config.llm_call_bounds() + overrides = {key: kwargs[key] for key in bounds if key in kwargs} + if overrides: + logging.info( + "llm_call label=%s bound overrides in effect: %s", label, overrides + ) + completion_args = { "model": custom_model if custom_model else Config.COMPLETION_MODEL, "messages": messages, "top_p": 1, - "timeout": Config.LLM_TIMEOUT, - # ``timeout`` is per attempt, so the retry budget has to be pinned too - # or the effective ceiling becomes a multiple of it. litellm's outer - # retry loop is disabled in favour of the SDK-level count. - "max_retries": Config.LLM_MAX_RETRIES, - "num_retries": 0, + **bounds, **kwargs, } diff --git a/api/config.py b/api/config.py index d51ccb0d..389e3521 100644 --- a/api/config.py +++ b/api/config.py @@ -49,11 +49,7 @@ def _embedding_kwargs(self) -> dict: thread runs it. ``timeout`` is per attempt, so the retry budget is pinned too or the effective ceiling becomes a multiple of it. """ - return { - "timeout": Config.LLM_TIMEOUT, - "max_retries": Config.LLM_MAX_RETRIES, - "num_retries": 0, - } + return Config.llm_call_bounds() def embed(self, text: Union[str, list]) -> list: """ @@ -224,6 +220,24 @@ class Config: # pylint: disable=too-many-instance-attributes # pylint: disable-next=invalid-name LLM_MAX_RETRIES: int = max(0, int(os.getenv("LLM_MAX_RETRIES", "1"))) + @classmethod + def llm_call_bounds(cls) -> dict: + """Provider kwargs that bound one logical LLM call end to end. + + ``LLM_TIMEOUT`` is the budget for the whole call, not per attempt. The + provider applies its timeout to each attempt, so the per-attempt value + is the budget divided by the number of attempts — otherwise a retry + pushes the real ceiling to a multiple of the configured one, which is + what made a 3s timeout take 10.8s to fail before the retry budget was + pinned. litellm's own retry loop stays off so the two cannot compound. + """ + attempts = cls.LLM_MAX_RETRIES + 1 + return { + "timeout": cls.LLM_TIMEOUT / attempts, + "max_retries": cls.LLM_MAX_RETRIES, + "num_retries": 0, + } + # Bounds for user-query execution against the target database. Offloading # execution to a thread stops a slow query from blocking other requests, # but nothing bounds how long the query itself runs without these. diff --git a/tests/test_embeddings_offloading.py b/tests/test_embeddings_offloading.py index ea943fed..2b23f49a 100644 --- a/tests/test_embeddings_offloading.py +++ b/tests/test_embeddings_offloading.py @@ -134,6 +134,11 @@ def test_async_callers_do_not_call_llms_inline(): def test_embedding_calls_are_time_bounded(): """A hung provider must not pin a worker thread forever.""" kwargs = Config.EMBEDDING_MODEL._embedding_kwargs() - assert kwargs["timeout"] == Config.LLM_TIMEOUT - assert kwargs["max_retries"] == Config.LLM_MAX_RETRIES + # LLM_TIMEOUT is the budget for the whole call, so the per-attempt value is + # that budget divided across attempts; retries cannot push the real ceiling + # past it. + attempts = Config.LLM_MAX_RETRIES + 1 + assert kwargs == Config.llm_call_bounds() + assert kwargs["timeout"] == Config.LLM_TIMEOUT / attempts + assert kwargs["timeout"] * attempts <= Config.LLM_TIMEOUT assert kwargs["num_retries"] == 0 diff --git a/tests/test_timeout_validation.py b/tests/test_timeout_validation.py index 99cc0090..4aa52bd7 100644 --- a/tests/test_timeout_validation.py +++ b/tests/test_timeout_validation.py @@ -45,3 +45,58 @@ def test_defaults_are_positive(monkeypatch): assert module.Config.DB_STATEMENT_TIMEOUT > 0 assert module.Config.LLM_TIMEOUT > 0 assert module.Config.LLM_MAX_RETRIES >= 0 + + +@pytest.mark.unit +@pytest.mark.parametrize("retries,expected_attempts", [(0, 1), (1, 2), (3, 4)]) +def test_llm_timeout_is_a_total_budget(monkeypatch, retries, expected_attempts): + """LLM_TIMEOUT bounds the whole call, not each attempt. + + The provider applies its timeout per attempt, so leaving the configured + value there made the real ceiling a multiple of it: a 3s timeout took + 10.8s to fail. The per-attempt value is now the budget divided by the + number of attempts. + """ + from api.config import Config + + monkeypatch.setattr(Config, "LLM_TIMEOUT", 90.0, raising=False) + monkeypatch.setattr(Config, "LLM_MAX_RETRIES", retries, raising=False) + + bounds = Config.llm_call_bounds() + assert bounds["max_retries"] == retries + # litellm's own retry loop stays off so the two cannot compound. + assert bounds["num_retries"] == 0 + assert bounds["timeout"] == 90.0 / expected_attempts + # The worst case across every attempt stays within the budget. + assert bounds["timeout"] * expected_attempts <= Config.LLM_TIMEOUT + + +@pytest.mark.unit +def test_embeddings_share_the_same_bounds(monkeypatch): + """One place defines the ceiling, so embeddings cannot drift from it.""" + from api.config import Config + + monkeypatch.setattr(Config, "LLM_TIMEOUT", 60.0, raising=False) + monkeypatch.setattr(Config, "LLM_MAX_RETRIES", 1, raising=False) + assert Config.EMBEDDING_MODEL._embedding_kwargs() == Config.llm_call_bounds() + + +@pytest.mark.unit +def test_call_site_bound_overrides_are_logged(monkeypatch, caplog): + """An explicit override is allowed but must not be silent.""" + import api.agents.utils as agent_utils + + def fake_completion(**kwargs): + assert kwargs["timeout"] == 1.5 + message = type("M", (), {"content": "ok"})() + choice = type("C", (), {"message": message})() + return type("R", (), {"choices": [choice]})() + + monkeypatch.setattr(agent_utils, "completion", fake_completion) + + with caplog.at_level("INFO"): + agent_utils.run_completion([{"role": "user", "content": "hi"}], + label="probe", timeout=1.5) + + assert "bound overrides in effect" in caplog.text, "override was not logged" + assert "timeout" in caplog.text From b1e29c05f86c6fb6a8431a425a9706e60f9af2e8 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 10:46:25 +0300 Subject: [PATCH 20/25] test: use one import style for api.config in the timeout tests The module was reached both ways in one file (`import api.config` for reload, `from api.config import Config` elsewhere). Reload needs the module object, so use an aliased `from api import config as api_config` and keep the file on a single style. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_timeout_validation.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_timeout_validation.py b/tests/test_timeout_validation.py index 4aa52bd7..d9f56c4a 100644 --- a/tests/test_timeout_validation.py +++ b/tests/test_timeout_validation.py @@ -13,8 +13,10 @@ def _reload_config(monkeypatch, **env): for key, value in env.items(): monkeypatch.setenv(key, value) - import api.config - return importlib.reload(api.config) + # ``from ... import`` throughout, for one consistent style; reload needs the + # module object, which the alias provides. + from api import config as api_config + return importlib.reload(api_config) @pytest.mark.unit From c3004a10d6ed0dfd12db6f923e60b19ffd768d76 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 12:06:32 +0300 Subject: [PATCH 21/25] fix(loaders): stop dropping URL options, bound socket reads, unbind the pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sixth review from @Naseem77. All four valid, and the first is a privilege bug I introduced earlier in this PR. **1. Introspection dropped URL connection options (security).** `options=` replaces the entire URL-supplied options string. `_execution_connect_kwargs` merges them, but when the schema deadline was added I wrote a raw `options=` at the introspection connect instead of reusing it — so `-c role=app_reader` was discarded and introspection connected as the URL's owning role, reading tables the connection had been scoped away from. His live probe showed exactly that: a restricted table extracted along with its sample value. The duplication is what allowed the divergence, so both paths now share `_connect_kwargs(url, budget_seconds)`; the only difference is the budget. **2. Deadlines did not bound socket reads.** A server-side `statement_timeout` only fires while the server is still talking to us; on a blackholed connection the client blocks in a read with no deadline. PostgreSQL connections now carry `tcp_user_timeout` (probed for libpq 12+, since older libpq rejects unknown keywords) plus keepalives, so the OS terminates the connection. Snowflake gets an explicit `socket_timeout` — `network_timeout` bounds retries, not reads, so a 5s configuration was still using the connector's 60s socket default. **3. The introspection cap broke across event loops.** A module-level `asyncio.Semaphore` binds to the first loop that contends on it and then raises `is bound to a different event loop`, which breaks any second `asyncio.run()`. Replaced with a dedicated `ThreadPoolExecutor`, which is loop-independent and bounds the worker threads themselves — so a cancelled introspection cannot free its slot while its worker is still running, which the semaphore needed a shield to approximate. **4. Valid PostgreSQL integer syntax was loosened.** The parser handled only decimal digits, so `077777` (octal, 32.767s), `0x10` (16ms) and `+5s` were replaced by the 60s ceiling — loosening stricter requests. It now implements the accepted grammar: optional sign, hex/octal/binary/decimal with digit separators, bare leading zero as octal, optional unit. Tests: 12 new cases. Both new guards were checked against the pre-fix behaviour — reinstating the raw `options=` fails with "URL role was dropped — introspection would run with more privilege", and reading a bare leading zero as decimal fails the octal case. The cross-loop test runs two separate `asyncio.run()` batches. 311 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/loaders/introspection.py | 81 ++++++++++++++-------------- api/loaders/postgres_loader.py | 75 +++++++++++++++++++++----- api/loaders/snowflake_loader.py | 7 +++ tests/test_db_execution_timeouts.py | 62 +++++++++++++++++++++ tests/test_schema_load_offloading.py | 62 ++++++++++++++++++++- 5 files changed, 230 insertions(+), 57 deletions(-) diff --git a/api/loaders/introspection.py b/api/loaders/introspection.py index 2a15e33e..a42c5b55 100644 --- a/api/loaders/introspection.py +++ b/api/loaders/introspection.py @@ -1,61 +1,60 @@ """Bounded execution for schema introspection. -Introspection runs in a worker thread because the drivers are blocking, and -cancelling the awaiting task does not stop that thread: a stalled database -holds its session and its worker until it answers. Two bounds follow from -that — a deadline the server applies (see each loader's connect parameters), -and a cap on how many introspections may occupy the shared executor at once, -so a stalled database cannot starve every other offloaded call. +Introspection runs off the event loop because the drivers are blocking, and +cancelling the awaiting task cannot stop the worker: a stalled database holds +its session and its worker until it answers. Two bounds follow — a deadline the +server applies (see each loader's connect parameters), and a cap on how many +introspections may run at once, so a stalled database cannot starve the shared +default executor that every other offloaded call uses. + +The cap is a dedicated ``ThreadPoolExecutor`` rather than an +``asyncio.Semaphore``: a module-level semaphore binds itself to the first loop +that contends on it and then raises ``is bound to a different event loop`` for +any later loop, and an executor bounds the *threads* themselves, so a cancelled +introspection cannot free its slot while its worker is still running. """ import asyncio +import functools import logging +import threading +from concurrent.futures import ThreadPoolExecutor from api.config import Config -_SLOTS: asyncio.Semaphore | None = None +_EXECUTOR: ThreadPoolExecutor | None = None +_EXECUTOR_LOCK = threading.Lock() -def _semaphore() -> asyncio.Semaphore: - """Create the semaphore lazily, on the loop that first needs it.""" - global _SLOTS # pylint: disable=global-statement - if _SLOTS is None: - _SLOTS = asyncio.Semaphore(Config.DB_SCHEMA_CONCURRENCY) - return _SLOTS +def _executor() -> ThreadPoolExecutor: + """Create the process-wide introspection pool once, lazily.""" + global _EXECUTOR # pylint: disable=global-statement + with _EXECUTOR_LOCK: + if _EXECUTOR is None: + _EXECUTOR = ThreadPoolExecutor( + max_workers=Config.DB_SCHEMA_CONCURRENCY, + thread_name_prefix="schema-introspect", + ) + return _EXECUTOR async def run_introspection(func, /, *args, **kwargs): - """Run *func* in a worker thread, holding one introspection slot. + """Run *func* on the bounded introspection pool. - The slot is released when the *thread* finishes, not when the awaiting task - ends. Cancelling the awaiting task cannot stop a running worker, so - releasing on cancellation would hand the slot to new work while the old - thread still holds a worker and a database session — letting the cap be - exceeded by exactly the disconnect-driven load it exists to bound. - - Cancellation still reaches the caller: the worker is shielded so it keeps - running (and keeps its slot) while the ``CancelledError`` propagates. + Cancellation reaches the caller while the worker keeps running: the future + is shielded, so an abandoned introspection continues to occupy its worker + until the driver returns. That is the point — releasing the slot early + would let the cap be exceeded by exactly the disconnect-driven load it + exists to bound. """ - slots = _semaphore() - if slots.locked(): + pool = _executor() + queued = getattr(pool, "_work_queue", None) + if queued is not None and queued.qsize(): logging.info( - "schema introspection queued: %d concurrent slots in use", + "schema introspection queued: all %d workers busy", Config.DB_SCHEMA_CONCURRENCY, ) - await slots.acquire() - try: - worker = asyncio.ensure_future(asyncio.to_thread(func, *args, **kwargs)) - except BaseException: - slots.release() - raise - - def _release(finished: asyncio.Future) -> None: - slots.release() - if not finished.cancelled(): - # Mark any failure retrieved: when the caller was cancelled nobody - # is left to await this future. - finished.exception() - - worker.add_done_callback(_release) - return await asyncio.shield(worker) + loop = asyncio.get_running_loop() + future = loop.run_in_executor(pool, functools.partial(func, *args, **kwargs)) + return await asyncio.shield(future) diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index 2052ec7d..38146426 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -18,6 +18,10 @@ # A real ``-c statement_timeout=`` directive, case-insensitive (GUC # names are), capturing an optionally quoted value that may carry a unit. +# ``tcp_user_timeout`` is a libpq 12+ connection parameter; older libpq +# rejects unknown keywords outright, so probe once rather than assume. +_TCP_USER_TIMEOUT_SUPPORTED = psycopg2.extensions.libpq_version() >= 120000 + _STATEMENT_TIMEOUT_RE = re.compile( # PostgreSQL accepts both directive forms in an options string: # ``-c name=value`` and ``--name=value``, the latter also with hyphens in @@ -157,26 +161,46 @@ def parse_schema_from_url(connection_url: str) -> str: def _statement_timeout_ms(raw: str): """Normalise a ``statement_timeout`` value to milliseconds. - PostgreSQL accepts a bare number (milliseconds) or a number with a unit - (``us``, ``ms``, ``s``, ``min``, ``h``, ``d``), optionally quoted. A - URL asking for ``5s`` is stricter than a 60s ceiling and should be - honoured, so the value has to be understood rather than pattern-matched - as digits. Returns ``None`` when the value is not something we can - compare, in which case the configured ceiling is used instead. + PostgreSQL accepts more than plain decimal digits: an optional sign, a + hexadecimal (``0x``), octal (``0o`` or a bare leading zero) or binary + (``0b``) integer, and an optional unit (``us``, ``ms``, ``s``, ``min``, + ``h``, ``d``), the whole thing possibly quoted. Only understanding + decimals meant `077777` (octal, 32.767s) and `0x10` (16ms) were treated + as unparseable or as decimals and then replaced by the ceiling — which + *loosened* a stricter request. Returns ``None`` when the value is not + something we can compare, in which case the ceiling applies. """ value = raw.strip().strip('"\'') match = re.fullmatch( - r"(?i)\s*(\d+(?:\.\d+)?)\s*(us|ms|s|min|h|d)?\s*", value + r"(?i)\s*([+-]?)\s*" # optional sign + r"(0[xX][0-9a-fA-F_]+|0[oO][0-7_]+|0[bB][01_]+|\d[\d_]*(?:\.\d+)?)" + r"\s*(us|ms|s|min|h|d)?\s*", # optional unit + value, ) if not match: return None - amount = float(match.group(1)) - unit = (match.group(2) or "ms").lower() + + sign, digits, unit = match.group(1), match.group(2).replace("_", ""), match.group(3) + try: + if digits[:2].lower() in ("0x", "0o", "0b"): + amount = float(int(digits, 0)) + elif "." in digits: + amount = float(digits) + elif len(digits) > 1 and digits.startswith("0"): + # A bare leading zero is octal to PostgreSQL, not decimal. + amount = float(int(digits, 8)) + else: + amount = float(digits) + except ValueError: + return None + + if sign == "-": + return None # negative disables nothing useful factors = { "us": 0.001, "ms": 1, "s": 1000, "min": 60_000, "h": 3_600_000, "d": 86_400_000, } - milliseconds = amount * factors[unit] + milliseconds = amount * factors[(unit or "ms").lower()] if milliseconds <= 0: return None # Round up so a sub-millisecond request (e.g. ``500us``) is honoured as @@ -207,10 +231,15 @@ def _introspect_schema(connection_url: str, schema: str): # than the user-query ceiling but still bounded: cancelling the # awaiting task cannot stop this thread, so a stalled database # would otherwise hold the session and the worker indefinitely. + # Shares the connect-keyword builder with query execution: a raw + # ``options=`` here would replace every URL-supplied option, and + # dropping something like ``-c role=app_reader`` would silently + # introspect with more privilege than the connection was granted. conn = psycopg2.connect( connection_url, - connect_timeout=Config.DB_CONNECT_TIMEOUT, - options=f"-c statement_timeout={Config.DB_SCHEMA_TIMEOUT * 1000}", + **PostgresLoader._connect_kwargs( + connection_url, Config.DB_SCHEMA_TIMEOUT + ), ) cursor = conn.cursor() @@ -602,7 +631,12 @@ async def refresh_graph_schema(graph_id: str, db_url: str, db=None) -> Tuple[boo @staticmethod def _execution_connect_kwargs(db_url: str) -> Dict[str, Any]: - """Timeout keywords for executing a user query. + """Connect keywords for executing a user query.""" + return PostgresLoader._connect_kwargs(db_url, Config.DB_STATEMENT_TIMEOUT) + + @staticmethod + def _connect_kwargs(db_url: str, statement_timeout_s: int) -> Dict[str, Any]: + """Connect keywords bounding one connection, in seconds. Offloading execution to a thread keeps the event loop free, but only a server-side ``statement_timeout`` bounds the query itself — a thread @@ -626,7 +660,7 @@ def _execution_connect_kwargs(db_url: str) -> Dict[str, Any]: # `statement_timeout=1000 ... statement_timeout=0` would end up # unbounded. GUC names are case-insensitive, so an uppercase directive # left behind would win over ours. - timeout_ms = Config.DB_STATEMENT_TIMEOUT * 1000 + timeout_ms = statement_timeout_s * 1000 requested = [ ms for ms in ( PostgresLoader._statement_timeout_ms(raw) @@ -652,6 +686,19 @@ def _execution_connect_kwargs(db_url: str) -> Dict[str, Any]: if 0 < requested <= connect_timeout: connect_timeout = requested kwargs["connect_timeout"] = connect_timeout + + # A server-side statement_timeout only fires while the server is still + # talking to us. On a blackholed connection — packets dropped rather + # than refused — the client blocks in a socket read with no deadline, + # keeping the worker alive well past the configured bound. TCP-level + # limits are what terminate that, so the OS gives up instead. + if "tcp_user_timeout" not in url_params and _TCP_USER_TIMEOUT_SUPPORTED: + kwargs["tcp_user_timeout"] = timeout_ms + if "keepalives" not in url_params: + kwargs["keepalives"] = 1 + kwargs["keepalives_idle"] = max(1, connect_timeout) + kwargs["keepalives_interval"] = max(1, connect_timeout) + kwargs["keepalives_count"] = 3 return kwargs @staticmethod diff --git a/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index 4adc0d65..4fa44881 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -261,6 +261,11 @@ def _introspect_schema( # the awaiting task cannot stop this thread. conn_params = dict(conn_params) conn_params["network_timeout"] = Config.DB_SCHEMA_TIMEOUT + # ``network_timeout`` bounds retries, not an individual socket + # read: without ``socket_timeout`` a stalled read falls back to + # the connector's own 60s default, ignoring the configured + # deadline entirely. + conn_params["socket_timeout"] = Config.DB_SCHEMA_TIMEOUT conn_params["session_parameters"] = { **(conn_params.get("session_parameters") or {}), "STATEMENT_TIMEOUT_IN_SECONDS": Config.DB_SCHEMA_TIMEOUT, @@ -682,6 +687,8 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: # configured values would never apply. conn_params["login_timeout"] = Config.DB_CONNECT_TIMEOUT conn_params["network_timeout"] = Config.DB_STATEMENT_TIMEOUT + # See the note in ``load``: retries are not socket reads. + conn_params["socket_timeout"] = Config.DB_STATEMENT_TIMEOUT session_parameters = dict(conn_params.get("session_parameters") or {}) session_parameters.setdefault( "STATEMENT_TIMEOUT_IN_SECONDS", Config.DB_STATEMENT_TIMEOUT diff --git a/tests/test_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py index 879fd2dc..c983d590 100644 --- a/tests/test_db_execution_timeouts.py +++ b/tests/test_db_execution_timeouts.py @@ -226,3 +226,65 @@ def test_postgres_clamp_keeps_unrelated_options(): f"-c statement_timeout={Config.DB_STATEMENT_TIMEOUT * 1000}" ) + + +@pytest.mark.unit +@pytest.mark.parametrize("budget_attr", ["DB_STATEMENT_TIMEOUT", "DB_SCHEMA_TIMEOUT"]) +def test_postgres_preserves_url_options_on_every_path(budget_attr): + """A privilege-bearing URL option must survive on both connect paths. + + ``options=`` replaces the whole URL-supplied options string, so building it + without merging drops things like ``-c role=app_reader`` — introspecting + with more privilege than the connection was granted. The schema path used a + raw ``options=`` and had exactly that bug. + """ + budget = getattr(Config, budget_attr) + kwargs = PostgresLoader._connect_kwargs( + f"{PG_URL}?options=-c%20role%3Dapp_reader%20-c%20search_path%3Drestricted", + budget, + ) + assert "role=app_reader" in kwargs["options"] + assert "search_path=restricted" in kwargs["options"] + assert f"-c statement_timeout={budget * 1000}" in kwargs["options"] + + +@pytest.mark.unit +def test_postgres_bounds_socket_reads_not_just_statements(): + """A server-side statement_timeout cannot fire on a blackholed socket. + + Packets dropped rather than refused leave the client blocked in a read with + no deadline, holding the worker past the configured bound, so the TCP-level + limits are what actually terminate it. + """ + kwargs = PostgresLoader._connect_kwargs(PG_URL, Config.DB_STATEMENT_TIMEOUT) + assert kwargs["keepalives"] == 1 + assert kwargs["keepalives_idle"] > 0 + assert kwargs["keepalives_count"] > 0 + # libpq 12+ only; the module probes support once at import. + from api.loaders.postgres_loader import _TCP_USER_TIMEOUT_SUPPORTED + if _TCP_USER_TIMEOUT_SUPPORTED: + assert kwargs["tcp_user_timeout"] == Config.DB_STATEMENT_TIMEOUT * 1000 + + +@pytest.mark.unit +@pytest.mark.parametrize("value,expected_ms,grammar", [ + ("077777", 32_767, "bare leading zero is octal"), + ("0o777", 511, "explicit octal"), + ("0x10", 16, "hexadecimal"), + ("0X1F", 31, "hexadecimal, uppercase"), + ("0b1010", 10, "binary"), + ("+5s", 5_000, "explicit positive sign"), + ("1_000", 1_000, "digit separators"), + ("-5", None, "negative is not a usable bound"), +]) +def test_postgres_parses_the_accepted_integer_grammar(value, expected_ms, grammar): + """PostgreSQL accepts more than decimal digits for an integer GUC. + + Reading `077777` as decimal 77777ms, or failing to parse `0x10` at all, + replaced a stricter request with the looser ceiling. + """ + ceiling = Config.DB_STATEMENT_TIMEOUT * 1000 + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options=-c%20statement_timeout%3D{value}" + ) + assert kwargs["options"] == f"-c statement_timeout={expected_ms or ceiling}", grammar diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py index b4b0fd42..0e056cb3 100644 --- a/tests/test_schema_load_offloading.py +++ b/tests/test_schema_load_offloading.py @@ -252,7 +252,7 @@ async def test_schema_introspection_concurrency_is_bounded(monkeypatch): """ import api.loaders.introspection as introspection - monkeypatch.setattr(introspection, "_SLOTS", None) + monkeypatch.setattr(introspection, "_EXECUTOR", None) monkeypatch.setattr(Config, "DB_SCHEMA_CONCURRENCY", 2, raising=False) live = 0 @@ -287,7 +287,7 @@ async def test_cancelled_introspection_keeps_its_slot(monkeypatch): """ import api.loaders.introspection as introspection - monkeypatch.setattr(introspection, "_SLOTS", None) + monkeypatch.setattr(introspection, "_EXECUTOR", None) monkeypatch.setattr(Config, "DB_SCHEMA_CONCURRENCY", 2, raising=False) live = 0 @@ -323,3 +323,61 @@ def blocking_work(): await asyncio.gather(*second, return_exceptions=True) + + +@pytest.mark.unit +def test_introspection_pool_survives_a_new_event_loop(monkeypatch): + """The cap must not be tied to whichever loop first contended on it. + + A module-level ``asyncio.Semaphore`` binds to the first loop that waits on + it and then raises ``is bound to a different event loop`` for every later + loop — which breaks any second `asyncio.run()`, including SDK callers. + """ + import api.loaders.introspection as introspection + + monkeypatch.setattr(introspection, "_EXECUTOR", None) + monkeypatch.setattr(Config, "DB_SCHEMA_CONCURRENCY", 2, raising=False) + + def work(): + time.sleep(0.05) + return "done" + + async def batch(): + # More work than workers, so the pool is genuinely contended. + return await asyncio.gather( + *(introspection.run_introspection(work) for _ in range(3)) + ) + + assert asyncio.run(batch()) == ["done"] * 3 + # A second, entirely separate loop must work just the same. + assert asyncio.run(batch()) == ["done"] * 3 + + +@pytest.mark.unit +@patch("api.loaders.postgres_loader.load_to_graph") +@patch("api.loaders.postgres_loader.psycopg2.connect") +async def test_introspection_connect_preserves_url_role(mock_connect, mock_load_to_graph): + """Introspection must not silently gain privilege the URL restricted. + + ``options=`` replaces the entire URL-supplied options string. Building it + without merging drops ``-c role=app_reader``, so introspection connects as + the URL's owning role and can read tables the connection was scoped away + from. This guards the connect call itself, not just the kwargs builder. + """ + async def noop(*_args, **_kwargs): + return None + + mock_load_to_graph.side_effect = noop + mock_connect.return_value = MagicMock() + + url = "postgresql://u:p@h:5432/db?options=-c%20role%3Dapp_reader" + with patch.object(PostgresLoader, "extract_tables_info", lambda *_a: {}), \ + patch.object(PostgresLoader, "extract_relationships", lambda *_a: {}): + async for _ in PostgresLoader.load("pfx", url): + pass + + options = mock_connect.call_args.kwargs["options"] + assert "role=app_reader" in options, ( + "URL role was dropped — introspection would run with more privilege" + ) + assert f"statement_timeout={Config.DB_SCHEMA_TIMEOUT * 1000}" in options From b071f913b9f58bd9975a9b3862426838fa2b930b Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 13:07:05 +0300 Subject: [PATCH 22/25] fix: drive retries against the remaining budget; enforce a real DB deadline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Seventh review from @Naseem77. All three valid; the first was a defect in the budget mechanism I added one review earlier. **1. The retry configuration disabled the retries it budgeted for.** litellm treats `num_retries` as overriding `max_retries`, so `{max_retries: 1, num_retries: 0}` made exactly one request while the timeout was divided as though two would happen — no retry at all, and half the intended deadline. Confirmed by counting requests against a local server: 1 where 2 were expected. Both library mechanisms are now off and `run_completion` owns the retry loop, handing each attempt what is left of the budget. Measured: retries=0 -> 1 request retries=1 -> 2 requests retries=3 -> 4 blackholed provider, 5s budget -> raised Timeout after 5.2s, 1 request A retry therefore happens only when time remains, which is the case worth retrying: a fast transient failure rather than a call that already spent the budget. Fixed a real bug found while testing this — the attempt kwargs were passed as several `**` expansions, so a caller overriding `timeout` hit `TypeError: got multiple values for keyword argument` instead of overriding. **2. TCP settings did not enforce the deadline.** Correct: `tcp_user_timeout` bounds unacknowledged outbound data and keepalives only detect a dead peer, so a stalled backend or proxy keeps TCP healthy while the client blocks in a read. Added `api/loaders/deadline.py` — a guard that cancels the statement at the deadline and closes the connection if the cancel does not take, which makes the blocked read raise and releases the worker. Applied to introspection and to user-query execution. URL socket settings are now clamped rather than deferred to, so `tcp_user_timeout=0&keepalives=0` can no longer switch the safeguards off; a URL may still tighten them. **3. Real-number timeout syntax was loosened.** `.5s` (500ms), `5.s` and `5e3ms` (5s) were all rewritten to the 60s ceiling. The grammar now accepts leading and trailing decimal points and exponent notation alongside the sign/radix forms. Tests: 20 new cases — attempt counts per retry setting, budget exhaustion stopping further attempts, library knobs staying off, the guard cancelling then closing (including when cancel itself fails), URL socket clamping, and the new numeric forms. Also fixed cross-file test pollution the reload-based validation tests were causing: they rebind `api.config.Config`, so tests patching a freshly imported reference were patching a different object than the module under test held. 333 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/agents/utils.py | 107 ++++++++++++++++++++-------- api/config.py | 32 +++++---- api/loaders/deadline.py | 71 ++++++++++++++++++ api/loaders/postgres_loader.py | 60 +++++++++++----- tests/test_db_execution_timeouts.py | 51 +++++++++++++ tests/test_deadline_guard.py | 77 ++++++++++++++++++++ tests/test_embeddings_offloading.py | 12 ++-- tests/test_timeout_validation.py | 105 +++++++++++++++++++++++---- 8 files changed, 434 insertions(+), 81 deletions(-) create mode 100644 api/loaders/deadline.py create mode 100644 tests/test_deadline_guard.py diff --git a/api/agents/utils.py b/api/agents/utils.py index 78da2951..6f46b57b 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -9,6 +9,35 @@ from api.config import Config +def _log_success(label: str, model: str, attempt: int, attempts: int, + elapsed: float) -> None: + """Record a completed call, flagging one slow enough to be worth noticing.""" + logging.info( + "llm_call label=%s model=%s attempt=%d/%d duration=%.2fs outcome=ok", + label, model, attempt, attempts, elapsed, + ) + if elapsed >= Config.LLM_SLOW_CALL_THRESHOLD: + logging.warning( + "llm_call label=%s model=%s duration=%.2fs exceeded slow-call " + "threshold of %.0fs", label, model, elapsed, + Config.LLM_SLOW_CALL_THRESHOLD, + ) + + +def _attempt(base_args: Dict[str, Any], remaining: float, overrides: Dict[str, Any]): + """Issue one provider request bounded by *remaining* seconds. + + Merged into one mapping rather than passed as several ``**`` expansions: + duplicate keywords are a ``TypeError`` that way, so a caller overriding + ``timeout`` would crash instead of overriding. + """ + return completion(**{ + **base_args, + **Config.llm_call_bounds(timeout=remaining), + **overrides, + }) + + def run_completion(messages: List[Dict[str, str]], custom_model: str | None = None, custom_api_key: str | None = None, *, label: str = "llm", **kwargs) -> str: @@ -26,45 +55,61 @@ def run_completion(messages: List[Dict[str, str]], custom_model: str | None = No Returns the content string from the first choice. """ - bounds = Config.llm_call_bounds() - overrides = {key: kwargs[key] for key in bounds if key in kwargs} - if overrides: - logging.info( - "llm_call label=%s bound overrides in effect: %s", label, overrides - ) - - completion_args = { + base_args = { "model": custom_model if custom_model else Config.COMPLETION_MODEL, "messages": messages, "top_p": 1, - **bounds, - **kwargs, } - if custom_api_key: - completion_args["api_key"] = custom_api_key + base_args["api_key"] = custom_api_key - started = time.monotonic() - try: - result = completion(**completion_args) - except Exception: - logging.warning( - "llm_call label=%s model=%s duration=%.2fs outcome=error", - label, completion_args["model"], time.monotonic() - started, + overrides = { + key: kwargs[key] + for key in ("timeout", "max_retries", "num_retries") + if key in kwargs + } + if overrides: + logging.info( + "llm_call label=%s bound overrides in effect: %s", label, overrides ) - raise - elapsed = time.monotonic() - started - logging.info( - "llm_call label=%s model=%s duration=%.2fs outcome=ok", - label, completion_args["model"], elapsed, + + attempts = Config.llm_attempts() + deadline = time.monotonic() + Config.LLM_TIMEOUT + last_error: Exception | None = None + + for attempt in range(1, attempts + 1): + remaining = deadline - time.monotonic() + if remaining <= 0: + break + + # Each attempt gets what is left of the budget, so the total cannot + # exceed LLM_TIMEOUT however many attempts are made. A retry therefore + # only happens when time remains — which is the case that matters, a + # fast transient failure rather than a call that already spent the + # budget. + started = time.monotonic() + try: + result = _attempt(base_args, remaining, kwargs) + except Exception as exc: # pylint: disable=broad-exception-caught + last_error = exc + logging.warning( + "llm_call label=%s model=%s attempt=%d/%d duration=%.2fs " + "outcome=error error=%s", + label, base_args["model"], attempt, attempts, + time.monotonic() - started, type(exc).__name__, + ) + continue + + _log_success(label, base_args["model"], attempt, attempts, + time.monotonic() - started) + return result.choices[0].message.content + + if last_error is not None: + raise last_error + raise TimeoutError( + f"llm_call label={label} exhausted its {Config.LLM_TIMEOUT}s budget " + "before an attempt could start" ) - if elapsed >= Config.LLM_SLOW_CALL_THRESHOLD: - logging.warning( - "llm_call label=%s model=%s duration=%.2fs exceeded slow-call " - "threshold of %.0fs", label, completion_args["model"], elapsed, - Config.LLM_SLOW_CALL_THRESHOLD, - ) - return result.choices[0].message.content class BaseAgent: # pylint: disable=too-few-public-methods diff --git a/api/config.py b/api/config.py index 389e3521..f4ab056d 100644 --- a/api/config.py +++ b/api/config.py @@ -221,20 +221,28 @@ class Config: # pylint: disable=too-many-instance-attributes LLM_MAX_RETRIES: int = max(0, int(os.getenv("LLM_MAX_RETRIES", "1"))) @classmethod - def llm_call_bounds(cls) -> dict: - """Provider kwargs that bound one logical LLM call end to end. - - ``LLM_TIMEOUT`` is the budget for the whole call, not per attempt. The - provider applies its timeout to each attempt, so the per-attempt value - is the budget divided by the number of attempts — otherwise a retry - pushes the real ceiling to a multiple of the configured one, which is - what made a 3s timeout take 10.8s to fail before the retry budget was - pinned. litellm's own retry loop stays off so the two cannot compound. + def llm_attempts(cls) -> int: + """How many provider requests one logical call may make.""" + return cls.LLM_MAX_RETRIES + 1 + + @classmethod + def llm_call_bounds(cls, timeout: float | None = None) -> dict: + """Provider kwargs for a single attempt. + + Both library retry mechanisms are disabled. litellm treats + ``num_retries`` as overriding ``max_retries``, so the pair + ``{max_retries: 1, num_retries: 0}`` made exactly one request while the + budget was divided as though two would happen — losing the retry and + halving the effective deadline at once. Retries are therefore driven by + ``run_completion`` against the remaining budget, where the attempt count + and elapsed time are observable. + + ``timeout`` defaults to the whole budget, which is right for + single-attempt callers such as embeddings. """ - attempts = cls.LLM_MAX_RETRIES + 1 return { - "timeout": cls.LLM_TIMEOUT / attempts, - "max_retries": cls.LLM_MAX_RETRIES, + "timeout": cls.LLM_TIMEOUT if timeout is None else timeout, + "max_retries": 0, "num_retries": 0, } diff --git a/api/loaders/deadline.py b/api/loaders/deadline.py new file mode 100644 index 00000000..c494c2df --- /dev/null +++ b/api/loaders/deadline.py @@ -0,0 +1,71 @@ +"""Application-level deadline for a blocking database connection. + +Driver and TCP settings do not cover every stall. A server-side +``statement_timeout`` only fires while the backend is processing our query; +``tcp_user_timeout`` bounds *unacknowledged outbound data*; keepalives only +detect a dead TCP peer. A stalled backend or a proxy that keeps the connection +alive without answering satisfies all three while the client stays blocked in a +read — holding a worker thread that cancellation cannot reclaim. + +This closes that gap from the outside: a timer cancels the in-flight query and, +failing that, closes the connection, which makes the blocked read raise in the +worker so the thread is released. +""" + +import contextlib +import logging +import threading + +# How long to wait for a cooperative cancel before closing the socket. +_CANCEL_GRACE_SECONDS = 5.0 + + +def _cancel(conn, label: str) -> None: + """Ask the server to abort the running statement, if the driver can.""" + cancel = getattr(conn, "cancel", None) + if cancel is None: + return + try: + cancel() + logging.warning("%s exceeded its deadline; cancelling the query", label) + except Exception as exc: # pylint: disable=broad-exception-caught + # PQcancel opens its own connection to the server, so it can fail or + # hang when the server is unreachable. The close below is the fallback. + logging.warning("%s cancel failed (%s); will close the connection", + label, type(exc).__name__) + + +def _close(conn, label: str) -> None: + """Force the socket shut so a blocked read raises instead of hanging.""" + try: + conn.close() + logging.warning("%s deadline exceeded; connection closed", label) + except Exception as exc: # pylint: disable=broad-exception-caught + logging.warning("%s close after deadline failed: %s", + label, type(exc).__name__) + + +@contextlib.contextmanager +def deadline_guard(conn, seconds: float, label: str = "database call"): + """Cancel, then close, *conn* if the body outlives *seconds*. + + The guard is what makes the configured deadline real for a connection whose + peer has stopped answering but has not dropped the socket. + """ + if not seconds or seconds <= 0: + yield + return + + cancel_timer = threading.Timer(seconds, _cancel, args=(conn, label)) + close_timer = threading.Timer( + seconds + _CANCEL_GRACE_SECONDS, _close, args=(conn, label) + ) + cancel_timer.daemon = True + close_timer.daemon = True + cancel_timer.start() + close_timer.start() + try: + yield + finally: + cancel_timer.cancel() + close_timer.cancel() diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index 38146426..8dc6529f 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -13,6 +13,7 @@ from api.config import Config from api.loaders.base_loader import BaseLoader # pylint: disable=import-error +from api.loaders.deadline import deadline_guard from api.loaders.graph_loader import load_to_graph # pylint: disable=import-error from api.loaders.introspection import run_introspection @@ -173,7 +174,10 @@ def _statement_timeout_ms(raw: str): value = raw.strip().strip('"\'') match = re.fullmatch( r"(?i)\s*([+-]?)\s*" # optional sign - r"(0[xX][0-9a-fA-F_]+|0[oO][0-7_]+|0[bB][01_]+|\d[\d_]*(?:\.\d+)?)" + r"(" + r"0[xX][0-9a-fA-F_]+|0[oO][0-7_]+|0[bB][01_]+" # radix-prefixed + r"|(?:\d[\d_]*\.?\d*|\.\d+)(?:[eE][+-]?\d+)?" # decimal, incl. + r")" # ".5", "5." and "5e3" r"\s*(us|ms|s|min|h|d)?\s*", # optional unit value, ) @@ -184,7 +188,8 @@ def _statement_timeout_ms(raw: str): try: if digits[:2].lower() in ("0x", "0o", "0b"): amount = float(int(digits, 0)) - elif "." in digits: + elif any(ch in digits for ch in ".eE"): + # Real syntax (".5", "5.", "5e3") is never octal. amount = float(digits) elif len(digits) > 1 and digits.startswith("0"): # A bare leading zero is octal to PostgreSQL, not decimal. @@ -243,14 +248,20 @@ def _introspect_schema(connection_url: str, schema: str): ) cursor = conn.cursor() - # Set the session search_path to the parsed schema so unqualified - # table references (e.g. in sample queries) resolve correctly. - cursor.execute( - sql.SQL("SET search_path TO {}").format(sql.Identifier(schema)) - ) - - entities = PostgresLoader.extract_tables_info(cursor, schema) - relationships = PostgresLoader.extract_relationships(cursor, schema) + # The server-side deadline only fires while the backend is + # answering; this one fires regardless, so a stalled peer cannot + # hold the worker past the budget. + with deadline_guard( + conn, Config.DB_SCHEMA_TIMEOUT, "schema introspection" + ): + # Set the session search_path to the parsed schema so + # unqualified table references resolve correctly. + cursor.execute( + sql.SQL("SET search_path TO {}").format(sql.Identifier(schema)) + ) + + entities = PostgresLoader.extract_tables_info(cursor, schema) + relationships = PostgresLoader.extract_relationships(cursor, schema) return entities, relationships finally: if cursor is not None: @@ -692,13 +703,24 @@ def _connect_kwargs(db_url: str, statement_timeout_s: int) -> Dict[str, Any]: # than refused — the client blocks in a socket read with no deadline, # keeping the worker alive well past the configured bound. TCP-level # limits are what terminate that, so the OS gives up instead. - if "tcp_user_timeout" not in url_params and _TCP_USER_TIMEOUT_SUPPORTED: - kwargs["tcp_user_timeout"] = timeout_ms - if "keepalives" not in url_params: - kwargs["keepalives"] = 1 - kwargs["keepalives_idle"] = max(1, connect_timeout) - kwargs["keepalives_interval"] = max(1, connect_timeout) - kwargs["keepalives_count"] = 3 + # Clamped, not defaulted: ``tcp_user_timeout=0&keepalives=0`` in a URL + # would otherwise switch these off entirely. A URL may tighten + # tcp_user_timeout, never loosen or disable it. + if _TCP_USER_TIMEOUT_SUPPORTED: + url_tcp = url_params.get("tcp_user_timeout", [None])[0] + tcp_user_timeout = timeout_ms + if url_tcp is not None: + try: + requested_ms = int(url_tcp) + except ValueError: + requested_ms = 0 + if 0 < requested_ms < tcp_user_timeout: + tcp_user_timeout = requested_ms + kwargs["tcp_user_timeout"] = tcp_user_timeout + kwargs["keepalives"] = 1 + kwargs["keepalives_idle"] = max(1, connect_timeout) + kwargs["keepalives_interval"] = max(1, connect_timeout) + kwargs["keepalives_count"] = 3 return kwargs @staticmethod @@ -719,9 +741,11 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: db_url, **PostgresLoader._execution_connect_kwargs(db_url) ) cursor = conn.cursor() + guard = deadline_guard(conn, Config.DB_STATEMENT_TIMEOUT, "query execution") # Execute the SQL query - cursor.execute(sql_query) + with guard: + cursor.execute(sql_query) # Check if the query returns results (SELECT queries) if cursor.description is not None: diff --git a/tests/test_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py index c983d590..45786bbe 100644 --- a/tests/test_db_execution_timeouts.py +++ b/tests/test_db_execution_timeouts.py @@ -288,3 +288,54 @@ def test_postgres_parses_the_accepted_integer_grammar(value, expected_ms, gramma f"{PG_URL}?options=-c%20statement_timeout%3D{value}" ) assert kwargs["options"] == f"-c statement_timeout={expected_ms or ceiling}", grammar + + +@pytest.mark.unit +@pytest.mark.parametrize("query,reason", [ + ("tcp_user_timeout=0&keepalives=0", "both disabled"), + ("keepalives=0", "keepalives disabled"), + ("tcp_user_timeout=0", "tcp_user_timeout disabled"), + ("tcp_user_timeout=999999999", "tcp_user_timeout loosened"), +]) +def test_postgres_url_cannot_disable_socket_safeguards(query, reason): + """A URL may tighten the socket bounds, never switch them off.""" + from api.loaders.postgres_loader import _TCP_USER_TIMEOUT_SUPPORTED + + kwargs = PostgresLoader._connect_kwargs( + f"{PG_URL}?{query}", Config.DB_STATEMENT_TIMEOUT + ) + assert kwargs["keepalives"] == 1, reason + if _TCP_USER_TIMEOUT_SUPPORTED: + assert kwargs["tcp_user_timeout"] == Config.DB_STATEMENT_TIMEOUT * 1000, reason + + +@pytest.mark.unit +def test_postgres_url_may_tighten_the_socket_bound(): + from api.loaders.postgres_loader import _TCP_USER_TIMEOUT_SUPPORTED + + kwargs = PostgresLoader._connect_kwargs( + f"{PG_URL}?tcp_user_timeout=5000", Config.DB_STATEMENT_TIMEOUT + ) + if _TCP_USER_TIMEOUT_SUPPORTED: + assert kwargs["tcp_user_timeout"] == 5000 + + +@pytest.mark.unit +@pytest.mark.parametrize("value,expected_ms,grammar", [ + (".5s", 500, "leading decimal point"), + ("5.s", 5_000, "trailing decimal point"), + ("5e3ms", 5_000, "exponent notation"), + ("5E3ms", 5_000, "exponent, uppercase"), + ("1e-1s", 100, "negative exponent"), + ("0.5s", 500, "ordinary decimal"), +]) +def test_postgres_parses_real_number_syntax(value, expected_ms, grammar): + """PostgreSQL accepts real syntax for a time GUC, not just integers. + + `.5s` is 500ms and `5e3ms` is 5s to the server; rejecting them replaced a + stricter request with the looser ceiling. + """ + kwargs = PostgresLoader._execution_connect_kwargs( + f"{PG_URL}?options=-c%20statement_timeout%3D{value.replace('+', '%2B')}" + ) + assert kwargs["options"] == f"-c statement_timeout={expected_ms}", grammar diff --git a/tests/test_deadline_guard.py b/tests/test_deadline_guard.py new file mode 100644 index 00000000..30c07982 --- /dev/null +++ b/tests/test_deadline_guard.py @@ -0,0 +1,77 @@ +"""The application-level deadline that driver settings cannot provide. + +A server-side ``statement_timeout`` only fires while the backend is answering, +``tcp_user_timeout`` bounds unacknowledged outbound data, and keepalives only +detect a dead TCP peer. A stalled backend or proxy satisfies all three while the +client stays blocked in a read, holding a worker thread that cancellation cannot +reclaim. The guard cancels and then closes, so the read raises. +""" + +import threading +import time + +import pytest + +import api.loaders.deadline as deadline +from api.loaders.deadline import deadline_guard + + +class _Conn: + """Records what the guard did to it.""" + + def __init__(self, cancel_raises=False): + self.cancelled = threading.Event() + self.closed = threading.Event() + self._cancel_raises = cancel_raises + + def cancel(self): + self.cancelled.set() + if self._cancel_raises: + raise OSError("server unreachable") + + def close(self): + self.closed.set() + + +@pytest.mark.unit +def test_guard_cancels_then_closes_a_stalled_call(monkeypatch): + monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) + conn = _Conn() + + with deadline_guard(conn, 0.1, "probe"): + # Stands in for a read that never returns. + assert conn.cancelled.wait(timeout=2), "deadline did not cancel the query" + assert conn.closed.wait(timeout=2), "deadline did not close the connection" + + +@pytest.mark.unit +def test_guard_closes_even_when_cancel_fails(monkeypatch): + """PQcancel opens its own connection, so it can fail on an unreachable server.""" + monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) + conn = _Conn(cancel_raises=True) + + with deadline_guard(conn, 0.1, "probe"): + assert conn.closed.wait(timeout=2), "close fallback did not run" + + +@pytest.mark.unit +def test_guard_leaves_a_prompt_call_alone(monkeypatch): + monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.1) + conn = _Conn() + + with deadline_guard(conn, 5.0, "probe"): + time.sleep(0.05) + + time.sleep(0.2) + assert not conn.cancelled.is_set(), "cancelled a call that finished in time" + assert not conn.closed.is_set(), "closed a call that finished in time" + + +@pytest.mark.unit +@pytest.mark.parametrize("seconds", [0, None, -1]) +def test_guard_is_a_noop_without_a_deadline(seconds): + conn = _Conn() + with deadline_guard(conn, seconds, "probe"): + pass + assert not conn.cancelled.is_set() + assert not conn.closed.is_set() diff --git a/tests/test_embeddings_offloading.py b/tests/test_embeddings_offloading.py index 2b23f49a..47318445 100644 --- a/tests/test_embeddings_offloading.py +++ b/tests/test_embeddings_offloading.py @@ -134,11 +134,11 @@ def test_async_callers_do_not_call_llms_inline(): def test_embedding_calls_are_time_bounded(): """A hung provider must not pin a worker thread forever.""" kwargs = Config.EMBEDDING_MODEL._embedding_kwargs() - # LLM_TIMEOUT is the budget for the whole call, so the per-attempt value is - # that budget divided across attempts; retries cannot push the real ceiling - # past it. - attempts = Config.LLM_MAX_RETRIES + 1 + # Embeddings are single-attempt, so they get the whole budget in one go. + # Library retries stay off here as they do on the completion path: litellm + # treats num_retries as overriding max_retries, which silently changed both + # the attempt count and the effective deadline. assert kwargs == Config.llm_call_bounds() - assert kwargs["timeout"] == Config.LLM_TIMEOUT / attempts - assert kwargs["timeout"] * attempts <= Config.LLM_TIMEOUT + assert kwargs["timeout"] == Config.LLM_TIMEOUT + assert kwargs["max_retries"] == 0 assert kwargs["num_retries"] == 0 diff --git a/tests/test_timeout_validation.py b/tests/test_timeout_validation.py index d9f56c4a..95665d56 100644 --- a/tests/test_timeout_validation.py +++ b/tests/test_timeout_validation.py @@ -6,10 +6,19 @@ """ import importlib +import time import pytest +@pytest.fixture(autouse=True) +def _restore_config_module(): + """Reload tests replace api.config; put the pristine module back after.""" + yield + from api import config as api_config + importlib.reload(api_config) + + def _reload_config(monkeypatch, **env): for key, value in env.items(): monkeypatch.setenv(key, value) @@ -52,25 +61,25 @@ def test_defaults_are_positive(monkeypatch): @pytest.mark.unit @pytest.mark.parametrize("retries,expected_attempts", [(0, 1), (1, 2), (3, 4)]) def test_llm_timeout_is_a_total_budget(monkeypatch, retries, expected_attempts): - """LLM_TIMEOUT bounds the whole call, not each attempt. + """LLM_TIMEOUT bounds the whole call, across every attempt. - The provider applies its timeout per attempt, so leaving the configured - value there made the real ceiling a multiple of it: a 3s timeout took - 10.8s to fail. The per-attempt value is now the budget divided by the - number of attempts. + The budget is enforced by the retry loop in ``run_completion``, which hands + each attempt the remaining time. The library's own retry knobs stay off: + litellm treats ``num_retries`` as overriding ``max_retries``, so relying on + them made one request while the budget was divided as though several would + happen. """ - from api.config import Config + from api import config as api_config - monkeypatch.setattr(Config, "LLM_TIMEOUT", 90.0, raising=False) - monkeypatch.setattr(Config, "LLM_MAX_RETRIES", retries, raising=False) + monkeypatch.setattr(api_config.Config, "LLM_TIMEOUT", 90.0, raising=False) + monkeypatch.setattr(api_config.Config, "LLM_MAX_RETRIES", retries, raising=False) - bounds = Config.llm_call_bounds() - assert bounds["max_retries"] == retries - # litellm's own retry loop stays off so the two cannot compound. + assert api_config.Config.llm_attempts() == expected_attempts + bounds = api_config.Config.llm_call_bounds() + assert bounds["max_retries"] == 0 assert bounds["num_retries"] == 0 - assert bounds["timeout"] == 90.0 / expected_attempts - # The worst case across every attempt stays within the budget. - assert bounds["timeout"] * expected_attempts <= Config.LLM_TIMEOUT + # Default is the whole budget, which is right for single-attempt callers. + assert bounds["timeout"] == 90.0 @pytest.mark.unit @@ -102,3 +111,71 @@ def fake_completion(**kwargs): assert "bound overrides in effect" in caplog.text, "override was not logged" assert "timeout" in caplog.text + + +@pytest.mark.unit +@pytest.mark.parametrize("retries", [0, 1, 3]) +def test_run_completion_makes_exactly_the_budgeted_attempts(monkeypatch, retries): + """Attempt count must match the configuration, and be observable. + + litellm treats ``num_retries`` as overriding ``max_retries``, so relying on + the library pair made one request while the budget was divided as though + several would happen — no retry, and half the deadline. Retries are driven + here instead. + """ + import api.agents.utils as agent_utils + + # Patch the Config the module under test holds: the reload-based tests in + # this file rebind api.config.Config, so a freshly imported reference can be + # a different object. + monkeypatch.setattr(agent_utils.Config, "LLM_TIMEOUT", 5.0, raising=False) + monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", retries, raising=False) + + calls = [] + + def failing_completion(**kwargs): + calls.append(kwargs["timeout"]) + raise RuntimeError("transient") + + monkeypatch.setattr(agent_utils, "completion", failing_completion) + + with pytest.raises(RuntimeError, match="transient"): + agent_utils.run_completion([{"role": "user", "content": "hi"}], label="probe") + + assert len(calls) == retries + 1 + # Library retries stay off, and each attempt is handed the remaining budget, + # so the per-attempt timeout never grows. + assert all(t <= 5.0 for t in calls) + assert calls == sorted(calls, reverse=True) + + +@pytest.mark.unit +def test_run_completion_stops_retrying_when_the_budget_is_spent(monkeypatch): + """A slow failure consumes the budget, so no further attempt is made.""" + import api.agents.utils as agent_utils + + monkeypatch.setattr(agent_utils.Config, "LLM_TIMEOUT", 0.3, raising=False) + monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", 5, raising=False) + + calls = [] + + def slow_failing_completion(**_kwargs): + calls.append(1) + time.sleep(0.2) + raise RuntimeError("slow transient") + + monkeypatch.setattr(agent_utils, "completion", slow_failing_completion) + + with pytest.raises(RuntimeError): + agent_utils.run_completion([{"role": "user", "content": "hi"}], label="probe") + + assert len(calls) < 6, "kept retrying past the budget" + + +@pytest.mark.unit +def test_library_retry_knobs_are_disabled(): + """Both library mechanisms stay off so they cannot compound or override.""" + import api.agents.utils as agent_utils + + bounds = agent_utils.Config.llm_call_bounds(timeout=7) + assert bounds == {"timeout": 7, "max_retries": 0, "num_retries": 0} From 817ee0e6ac860b7c9db6a15c016f6258a7868f8b Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 13:14:08 +0300 Subject: [PATCH 23/25] test: use one import style for the deadline module The guard tests monkeypatch the module's grace constant, so the module object is needed regardless; drop the parallel `from ... import deadline_guard` and call through the module. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_deadline_guard.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/test_deadline_guard.py b/tests/test_deadline_guard.py index 30c07982..24841ef9 100644 --- a/tests/test_deadline_guard.py +++ b/tests/test_deadline_guard.py @@ -12,8 +12,9 @@ import pytest -import api.loaders.deadline as deadline -from api.loaders.deadline import deadline_guard +# One import style: the module object is needed anyway, because the tests +# monkeypatch its grace constant. +from api.loaders import deadline class _Conn: @@ -38,7 +39,7 @@ def test_guard_cancels_then_closes_a_stalled_call(monkeypatch): monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) conn = _Conn() - with deadline_guard(conn, 0.1, "probe"): + with deadline.deadline_guard(conn, 0.1, "probe"): # Stands in for a read that never returns. assert conn.cancelled.wait(timeout=2), "deadline did not cancel the query" assert conn.closed.wait(timeout=2), "deadline did not close the connection" @@ -50,7 +51,7 @@ def test_guard_closes_even_when_cancel_fails(monkeypatch): monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) conn = _Conn(cancel_raises=True) - with deadline_guard(conn, 0.1, "probe"): + with deadline.deadline_guard(conn, 0.1, "probe"): assert conn.closed.wait(timeout=2), "close fallback did not run" @@ -59,7 +60,7 @@ def test_guard_leaves_a_prompt_call_alone(monkeypatch): monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.1) conn = _Conn() - with deadline_guard(conn, 5.0, "probe"): + with deadline.deadline_guard(conn, 5.0, "probe"): time.sleep(0.05) time.sleep(0.2) @@ -71,7 +72,7 @@ def test_guard_leaves_a_prompt_call_alone(monkeypatch): @pytest.mark.parametrize("seconds", [0, None, -1]) def test_guard_is_a_noop_without_a_deadline(seconds): conn = _Conn() - with deadline_guard(conn, seconds, "probe"): + with deadline.deadline_guard(conn, seconds, "probe"): pass assert not conn.cancelled.is_set() assert not conn.closed.is_set() From 6d8dab3fdb6b6a724e7bccfb199f33a82e42ec11 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 15:10:56 +0300 Subject: [PATCH 24/25] fix: unblockable DB deadline, guard through commit, verdict-based LLM retries Three fixes from review: The deadline guard escalated through conn.close(), but psycopg2 holds a connection lock during a blocking read, so the close from the timer thread joined the deadlock instead of breaking it - and PQcancel can hang against a black-holed server with nothing bounding it. The guard now duplicates the connection's socket at entry and escalates by shutdown(2) on it: no driver lock, no network round trip, and the blocked read raises at once. The dup is taken at entry (never at fire time) so a recycled descriptor number can never be hit, and closing it on exit disarms a late callback. Reproduced in tests against a real blocked socket read holding the driver lock, with a hanging cancel. The postgres guard ended right after cursor.execute(), leaving fetch, commit and rollback unbounded - a peer that answers the query but stalls on the commit held the worker just as effectively. One guard now spans all of them, and loader cleanup is suppressed so a rollback/close failure cannot mask the error that got there (the mysql and snowflake loaders had the same masking pattern, plus an unbound-cursor NameError in their error paths). LLM retries now run on one verdict: transport failures and 408/429/5xx are transient and retried against the remaining budget with backoff (Retry-After honoured, capped); 4xx verdicts about the request itself fail immediately - a 401 does not become a valid key by asking twice. run_batch_completion applies the same policy to batch calls, retrying only the failed slots, which replaces the conflicting {max_retries, num_retries} pair that litellm resolved to no retry at all; the key-validation probe in settings uses the shared single- attempt bounds for the same reason. Co-Authored-By: Claude Fable 5 --- api/agents/utils.py | 160 +++++++++++++++++++++++-- api/loaders/deadline.py | 126 ++++++++++++++++--- api/loaders/mysql_loader.py | 32 ++--- api/loaders/postgres_loader.py | 128 +++++++++++--------- api/loaders/snowflake_loader.py | 32 ++--- api/routes/settings.py | 7 +- api/utils.py | 18 ++- tests/test_deadline_guard.py | 141 +++++++++++++++++++++- tests/test_llm_retry_policy.py | 200 +++++++++++++++++++++++++++++++ tests/test_settings_route.py | 12 +- tests/test_timeout_validation.py | 5 + 11 files changed, 733 insertions(+), 128 deletions(-) create mode 100644 tests/test_llm_retry_policy.py diff --git a/api/agents/utils.py b/api/agents/utils.py index 6f46b57b..62eca64c 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -5,9 +5,15 @@ import time from typing import Any, Dict, List -from litellm import completion +from litellm import batch_completion, completion from api.config import Config +# Backoff between retryable failures: base doubles per attempt, capped so a +# late retry cannot sleep away the whole budget. Module-level so tests can +# shrink it. +_RETRY_BACKOFF_SECONDS = 0.5 +_RETRY_BACKOFF_CAP_SECONDS = 8.0 + def _log_success(label: str, model: str, attempt: int, attempts: int, elapsed: float) -> None: @@ -24,6 +30,71 @@ def _log_success(label: str, model: str, attempt: int, attempts: int, ) +def _retryable(exc: Exception) -> bool: + """Whether a failed attempt is worth spending remaining budget on. + + Transport-level failures (timeouts, dropped connections) carry no status + code and are treated as transient. When the provider did answer, its + verdict decides: 408/429/5xx describe the service's moment and can change + on a replay; everything else (400, 401, 403, 404, ...) describes the + request itself, so a retry would replay the same failure — a bad API key + does not become valid by asking twice. + """ + status_code = getattr(exc, "status_code", None) + if status_code is None: + return True + return status_code in (408, 429) or status_code >= 500 + + +def _retry_after_seconds(exc: Exception) -> float | None: + """The provider's Retry-After, if the failure carried one.""" + headers = getattr(getattr(exc, "response", None), "headers", None) + if headers is None: + return None + try: + value = headers.get("retry-after") + return float(value) if value is not None else None + except (TypeError, ValueError): + return None + + +def _retry_delay(exc: Exception, attempt: int) -> float: + """Seconds to wait before the next attempt. + + The provider's own Retry-After wins when present (a 429 tells us exactly + when trying again stops being rude); otherwise a small exponential + backoff, so a struggling service is not hammered at full speed. + """ + retry_after = _retry_after_seconds(exc) + if retry_after is not None and retry_after >= 0: + return min(retry_after, _RETRY_BACKOFF_CAP_SECONDS) + return min(_RETRY_BACKOFF_SECONDS * (2 ** (attempt - 1)), + _RETRY_BACKOFF_CAP_SECONDS) + + +def _pause_before_retry(delay: float, deadline: float) -> bool: + """Sleep *delay* before the next attempt. + + Returns ``False`` when sleeping would outlive the budget, in which case + there is no point in another attempt at all. + """ + if delay >= deadline - time.monotonic(): + return False + time.sleep(delay) + return True + + +def _log_failure(label: str, model: str, attempt: str, *, + elapsed: float, exc: Exception) -> None: + """Record a failed attempt and whether it is worth replaying.""" + logging.warning( + "llm_call label=%s model=%s attempt=%s duration=%.2fs " + "outcome=error error=%s retryable=%s", + label, model, attempt, elapsed, + type(exc).__name__, _retryable(exc), + ) + + def _attempt(base_args: Dict[str, Any], remaining: float, overrides: Dict[str, Any]): """Issue one provider request bounded by *remaining* seconds. @@ -92,12 +163,15 @@ def run_completion(messages: List[Dict[str, str]], custom_model: str | None = No result = _attempt(base_args, remaining, kwargs) except Exception as exc: # pylint: disable=broad-exception-caught last_error = exc - logging.warning( - "llm_call label=%s model=%s attempt=%d/%d duration=%.2fs " - "outcome=error error=%s", - label, base_args["model"], attempt, attempts, - time.monotonic() - started, type(exc).__name__, - ) + _log_failure(label, base_args["model"], f"{attempt}/{attempts}", + elapsed=time.monotonic() - started, exc=exc) + if not _retryable(exc): + # The provider judged the request itself invalid; a replay + # would fail identically, so surface it now. + raise + if attempt < attempts and not _pause_before_retry( + _retry_delay(exc, attempt), deadline): + break # sleeping would outlive the budget continue _log_success(label, base_args["model"], attempt, attempts, @@ -112,6 +186,78 @@ def run_completion(messages: List[Dict[str, str]], custom_model: str | None = No ) +def _fail_unanswered_slots(results: List[Any], label: str) -> None: + """Turn every still-``None`` slot into an explicit failure. + + A slot can hold ``None`` when the library answered short or the budget ran + out before the slot was ever attempted. Callers branch on + ``isinstance(..., Exception)`` and treat everything else as a response, so + ``None`` must leave as a failure, not a response. + """ + for i, item in enumerate(results): + if item is None: + results[i] = TimeoutError( + f"llm_call label={label} slot {i} got no answer within the " + f"{Config.LLM_TIMEOUT}s budget" + ) + + +def run_batch_completion(messages_list: List[List[Dict[str, str]]], *, + label: str = "llm-batch", **base_args) -> List[Any]: + """Batch counterpart of ``run_completion``: same budget, same verdicts. + + litellm's ``batch_completion`` reports a failed item by returning the + exception in that item's slot, and its own retry knobs conflict (it treats + ``num_retries`` as overriding ``max_retries``), so callers passing the pair + got no retry at all. This drives the whole batch against one ``LLM_TIMEOUT`` + budget instead and retries only the items whose failure was transient + (:func:`_retryable`), with whatever budget remains. + + Returns a list aligned with *messages_list*; an item that never succeeded + holds its final exception, which is the contract callers already handle. + """ + results: List[Any] = [None] * len(messages_list) + pending = list(range(len(messages_list))) + attempts = Config.llm_attempts() + deadline = time.monotonic() + Config.LLM_TIMEOUT + + for attempt in range(1, attempts + 1): + remaining = deadline - time.monotonic() + if not pending or remaining <= 0: + break + + started = time.monotonic() + batch = batch_completion(**{ + **base_args, + "messages": [messages_list[i] for i in pending], + **Config.llm_call_bounds(timeout=remaining), + }) + + retry_slots = [] + for slot, response in zip(pending, batch): + results[slot] = response + if isinstance(response, Exception) and _retryable(response): + retry_slots.append(slot) + + failed = sum(1 for i in pending if isinstance(results[i], Exception)) + logging.info( + "llm_call label=%s model=%s attempt=%d/%d duration=%.2fs " + "outcome=%d/%d ok (%d retryable)", + label, base_args.get("model"), attempt, attempts, + time.monotonic() - started, len(pending) - failed, len(pending), + len(retry_slots), + ) + + pending = retry_slots + if pending and attempt < attempts and not _pause_before_retry( + max(_retry_delay(results[i], attempt) for i in pending), + deadline): + break # sleeping would outlive the budget + + _fail_unanswered_slots(results, label) + return results + + class BaseAgent: # pylint: disable=too-few-public-methods """Base class for agents.""" diff --git a/api/loaders/deadline.py b/api/loaders/deadline.py index c494c2df..e47bc192 100644 --- a/api/loaders/deadline.py +++ b/api/loaders/deadline.py @@ -7,36 +7,109 @@ alive without answering satisfies all three while the client stays blocked in a read — holding a worker thread that cancellation cannot reclaim. -This closes that gap from the outside: a timer cancels the in-flight query and, -failing that, closes the connection, which makes the blocked read raise in the -worker so the thread is released. +This closes that gap from the outside: a timer asks the server to cancel the +in-flight query and, failing that, shuts the connection's socket down at the OS +level, which makes the blocked read raise in the worker so the thread is +released. + +The escalation must not itself be blockable by the stall it exists to break: + +* ``conn.cancel()`` (PQcancel) opens its *own* connection to the server, so + against a black-holed host the cancel itself can hang. It therefore runs on + a dedicated daemon timer thread that nothing waits for — a hanging cancel + costs one parked thread until the kernel gives up, never the escalation. +* ``conn.close()`` is not usable from another thread while a query is running: + psycopg2 serialises connection access with an internal lock that the worker + blocked in ``recv`` is holding, so a close from the timer thread would join + the deadlock instead of breaking it. The escalation instead calls + ``shutdown(2)`` on the underlying socket via a duplicated file descriptor — + a plain syscall that needs no driver lock and no network round trip. The + kernel fails the pending read immediately, the driver raises in the worker, + and the connection is then closed by the worker's own cleanup path, which + holds the lock legitimately. """ import contextlib import logging +import os +import socket import threading -# How long to wait for a cooperative cancel before closing the socket. +# How long to wait for a cooperative cancel before shutting the socket down. _CANCEL_GRACE_SECONDS = 5.0 def _cancel(conn, label: str) -> None: - """Ask the server to abort the running statement, if the driver can.""" + """Ask the server to abort the running statement, if the driver can. + + Best-effort only: this may hang (see the module docstring), and the + escalation to ``_shutdown`` does not depend on it returning. + """ cancel = getattr(conn, "cancel", None) if cancel is None: return + logging.warning("%s exceeded its deadline; cancelling the query", label) try: cancel() - logging.warning("%s exceeded its deadline; cancelling the query", label) except Exception as exc: # pylint: disable=broad-exception-caught - # PQcancel opens its own connection to the server, so it can fail or - # hang when the server is unreachable. The close below is the fallback. - logging.warning("%s cancel failed (%s); will close the connection", + # PQcancel opens its own connection to the server, so it can fail when + # the server is unreachable. The socket shutdown is the fallback. + logging.warning("%s cancel failed (%s); the socket shutdown will follow", + label, type(exc).__name__) + + +def _connection_fd(conn): + """The connection's socket descriptor, or ``None`` if it has none.""" + fileno = getattr(conn, "fileno", None) + if fileno is None: + return None + try: + fd = fileno() + except Exception: # pylint: disable=broad-exception-caught + # e.g. psycopg2.InterfaceError once the connection is already closed. + return None + return fd if isinstance(fd, int) and fd >= 0 else None + + +def _duplicate_socket(conn): + """A second handle on the connection's socket, or ``None``. + + Taken while the caller still demonstrably owns the connection, so the + deadline callback never resolves a file descriptor at fire time — a + descriptor number can be closed and reused by an unrelated file between + those two moments, and a shutdown aimed by number could then hit it. The + duplicate is an object handle on the *socket itself*: after ``close`` it + refuses further use instead of chasing a recycled number. + """ + fd = _connection_fd(conn) + if fd is None: + return None + try: + return socket.socket(fileno=os.dup(fd)) + except OSError: + return None + + +def _shutdown(dup: socket.socket, label: str) -> None: + """Make the blocked read raise, without asking the driver for anything. + + ``shutdown`` acts on the socket, which the driver's descriptor and this + duplicate share, so the worker's pending ``recv`` fails at once — no + driver lock, no network round trip, nothing on this path can block. + """ + try: + dup.shutdown(socket.SHUT_RDWR) + logging.warning("%s deadline exceeded; socket shut down so the blocked " + "call raises", label) + except OSError as exc: + # ENOTCONN and the like: the peer (or the guard's own exit, which + # closes the duplicate) beat us to it — the outcome we wanted anyway. + logging.warning("%s socket shutdown after deadline was a no-op: %s", label, type(exc).__name__) def _close(conn, label: str) -> None: - """Force the socket shut so a blocked read raises instead of hanging.""" + """Driver-level close, for connections that expose no socket.""" try: conn.close() logging.warning("%s deadline exceeded; connection closed", label) @@ -47,25 +120,44 @@ def _close(conn, label: str) -> None: @contextlib.contextmanager def deadline_guard(conn, seconds: float, label: str = "database call"): - """Cancel, then close, *conn* if the body outlives *seconds*. + """Cancel, then shut down, *conn* if the body outlives *seconds*. The guard is what makes the configured deadline real for a connection whose - peer has stopped answering but has not dropped the socket. + peer has stopped answering but has not dropped the socket. Both escalation + steps run on their own daemon timer threads and neither waits for the + other, so a step that itself hangs cannot postpone the next one. """ if not seconds or seconds <= 0: yield return + dup = _duplicate_socket(conn) + if dup is not None: + # The normal case: escalate by shutting the shared socket down. + escalate, escalate_args = _shutdown, (dup, label) + else: + # No reachable socket (a driver that exposes none, or a test fake): + # the driver-level close is all that is left. It can block on the + # driver's connection lock, which is exactly why the socket path is + # preferred whenever it exists. + escalate, escalate_args = _close, (conn, label) + cancel_timer = threading.Timer(seconds, _cancel, args=(conn, label)) - close_timer = threading.Timer( - seconds + _CANCEL_GRACE_SECONDS, _close, args=(conn, label) + shutdown_timer = threading.Timer( + seconds + _CANCEL_GRACE_SECONDS, escalate, args=escalate_args ) cancel_timer.daemon = True - close_timer.daemon = True + shutdown_timer.daemon = True cancel_timer.start() - close_timer.start() + shutdown_timer.start() try: yield finally: cancel_timer.cancel() - close_timer.cancel() + shutdown_timer.cancel() + if dup is not None: + # Closing the duplicate also disarms a shutdown callback that + # already started: the socket object refuses use after close, so + # a late timer cannot touch whatever the kernel reuses the + # descriptor number for. + dup.close() diff --git a/api/loaders/mysql_loader.py b/api/loaders/mysql_loader.py index fdeaeb19..7b4f5a8c 100644 --- a/api/loaders/mysql_loader.py +++ b/api/loaders/mysql_loader.py @@ -1,5 +1,6 @@ """MySQL loader for loading database schemas into FalkorDB graphs.""" +import contextlib import datetime import decimal import logging @@ -533,6 +534,8 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: Returns: List of dictionaries containing the query results """ + conn = None + cursor = None try: # Parse connection URL conn_params = MySQLLoader._parse_mysql_url(db_url) @@ -586,25 +589,26 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: # Commit the transaction for write operations conn.commit() - # Close database connection - cursor.close() - conn.close() - return result_list except pymysql.MySQLError as e: - # Rollback in case of error - if 'conn' in locals(): - conn.rollback() - cursor.close() - conn.close() + # Bounded by the socket timeouts above, and suppressed so a + # rollback failure cannot mask the error that got us here. + if conn is not None: + with contextlib.suppress(Exception): + conn.rollback() logging.error("MySQL query execution error: %s", e) raise MySQLQueryError(f"MySQL query execution error: {str(e)}") from e except Exception as e: - # Rollback in case of error - if 'conn' in locals(): - conn.rollback() - cursor.close() - conn.close() + if conn is not None: + with contextlib.suppress(Exception): + conn.rollback() logging.error("Error executing SQL query: %s", e) raise MySQLQueryError(f"Error executing SQL query: {str(e)}") from e + finally: + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() + if conn is not None: + with contextlib.suppress(Exception): + conn.close() diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index 8dc6529f..6858aff9 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -1,5 +1,6 @@ """PostgreSQL loader for loading database schemas into FalkorDB graphs.""" +import contextlib import re import datetime import decimal @@ -264,10 +265,14 @@ def _introspect_schema(connection_url: str, schema: str): relationships = PostgresLoader.extract_relationships(cursor, schema) return entities, relationships finally: + # A raise from cleanup (e.g. on a socket the deadline guard shut + # down) must not mask the introspection error itself. if cursor is not None: - cursor.close() + with contextlib.suppress(Exception): + cursor.close() if conn is not None: - conn.close() + with contextlib.suppress(Exception): + conn.close() @staticmethod async def load( # pylint: disable=arguments-differ @@ -736,69 +741,80 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: Returns: List of dictionaries containing the query results """ + conn = None + cursor = None try: conn = psycopg2.connect( db_url, **PostgresLoader._execution_connect_kwargs(db_url) ) cursor = conn.cursor() - guard = deadline_guard(conn, Config.DB_STATEMENT_TIMEOUT, "query execution") - - # Execute the SQL query - with guard: - cursor.execute(sql_query) - - # Check if the query returns results (SELECT queries) - if cursor.description is not None: - # This is a SELECT query or similar that returns rows - columns = [desc[0] for desc in cursor.description] - results = cursor.fetchall() - result_list = [] - for row in results: - # Serialize each value to ensure JSON compatibility - serialized_row = { - columns[i]: PostgresLoader._serialize_value(row[i]) - for i in range(len(columns)) - } - result_list.append(serialized_row) - else: - # This is an INSERT, UPDATE, DELETE, or other non-SELECT query - # Return information about the operation - affected_rows = cursor.rowcount - sql_type = sql_query.strip().split()[0].upper() - - if sql_type in ['INSERT', 'UPDATE', 'DELETE']: - result_list = [{ - "operation": sql_type, - "affected_rows": affected_rows, - "status": "success" - }] - else: - # For other types of queries (CREATE, DROP, etc.) - result_list = [{ - "operation": sql_type, - "status": "success" - }] - - # Commit the transaction for write operations - conn.commit() - - # Close database connection - cursor.close() - conn.close() + + # One guard spans execute, fetch, commit and rollback: a peer that + # answers the query but stalls while returning rows or on the + # commit holds the worker just as effectively as one that stalls + # on the execute itself. + with deadline_guard( + conn, Config.DB_STATEMENT_TIMEOUT, "query execution" + ): + try: + cursor.execute(sql_query) + + # Check if the query returns results (SELECT queries) + if cursor.description is not None: + # This is a SELECT query or similar that returns rows + columns = [desc[0] for desc in cursor.description] + results = cursor.fetchall() + result_list = [] + for row in results: + # Serialize each value to ensure JSON compatibility + serialized_row = { + columns[i]: PostgresLoader._serialize_value(row[i]) + for i in range(len(columns)) + } + result_list.append(serialized_row) + else: + # This is an INSERT, UPDATE, DELETE, or other + # non-SELECT query - return information about it + affected_rows = cursor.rowcount + sql_type = sql_query.strip().split()[0].upper() + + if sql_type in ['INSERT', 'UPDATE', 'DELETE']: + result_list = [{ + "operation": sql_type, + "affected_rows": affected_rows, + "status": "success" + }] + else: + # For other types of queries (CREATE, DROP, etc.) + result_list = [{ + "operation": sql_type, + "status": "success" + }] + + # Commit the transaction for write operations + conn.commit() + except Exception: + # The rollback is a network round trip too, so it stays + # inside the guard - and it is suppressed so it cannot + # mask the error that got us here (after a deadline + # shutdown it raises immediately on the dead socket). + with contextlib.suppress(Exception): + conn.rollback() + raise return result_list except psycopg2.Error as e: - # Rollback in case of error - if 'conn' in locals(): - conn.rollback() - cursor.close() - conn.close() raise PostgreSQLConnectionError(f"PostgreSQL query execution error: {str(e)}") from e except Exception as e: - # Rollback in case of error - if 'conn' in locals(): - conn.rollback() - cursor.close() - conn.close() raise PostgreSQLQueryError(f"Error executing SQL query: {str(e)}") from e + finally: + # Runs on a healthy connection or on one whose socket the guard + # already shut down; either way a raise from cleanup must not + # mask the real error. + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() + if conn is not None: + with contextlib.suppress(Exception): + conn.close() diff --git a/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index 4fa44881..d60a3b33 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -1,6 +1,7 @@ """Snowflake loader for loading database schemas into FalkorDB graphs.""" import base64 +import contextlib import datetime import decimal import logging @@ -677,6 +678,8 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: Returns: List of dictionaries containing the query results """ + conn = None + cursor = None try: # Parse connection URL conn_params = SnowflakeLoader._parse_snowflake_url(db_url) @@ -736,25 +739,26 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: # Commit the transaction for write operations conn.commit() - # Close database connection - cursor.close() - conn.close() - return result_list except snowflake.connector.Error as e: - # Rollback in case of error - if 'conn' in locals(): - conn.rollback() - cursor.close() - conn.close() + # Bounded by the network/socket timeouts above, and suppressed so + # a rollback failure cannot mask the error that got us here. + if conn is not None: + with contextlib.suppress(Exception): + conn.rollback() logging.error("Snowflake query execution error: %s", e) raise SnowflakeQueryError(f"Snowflake query execution error: {str(e)}") from e except Exception as e: - # Rollback in case of error - if 'conn' in locals(): - conn.rollback() - cursor.close() - conn.close() + if conn is not None: + with contextlib.suppress(Exception): + conn.rollback() logging.error("Error executing SQL query: %s", e) raise SnowflakeQueryError(f"Error executing SQL query: {str(e)}") from e + finally: + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() + if conn is not None: + with contextlib.suppress(Exception): + conn.close() diff --git a/api/routes/settings.py b/api/routes/settings.py index 2e77c49c..07598811 100644 --- a/api/routes/settings.py +++ b/api/routes/settings.py @@ -87,9 +87,10 @@ async def validate_api_key(request: Request, data: ValidateKeyRequest): # pylin messages=[{"role": "user", "content": "test"}], max_tokens=1, api_key=api_key, - timeout=Config.LLM_TIMEOUT, - max_retries=Config.LLM_MAX_RETRIES, - num_retries=0, + # One bounded attempt, retries off: a key-validation probe + # wants the provider's verdict, and a 401 is that verdict - + # replaying it would only delay the answer. + **Config.llm_call_bounds(), ) ) diff --git a/api/utils.py b/api/utils.py index 9fd7c6fc..2e4a72b3 100644 --- a/api/utils.py +++ b/api/utils.py @@ -2,9 +2,7 @@ import json from typing import Dict, List, Optional, TypedDict -from litellm import batch_completion - -from api.agents.utils import run_completion +from api.agents.utils import run_batch_completion, run_completion from api.config import Config @@ -84,16 +82,16 @@ def create_combined_description( # pylint: disable=too-many-locals for batch_start in range(0, len(messages_list), batch_size): batch_messages = messages_list[batch_start : batch_start + batch_size] - # Bounded like every other provider call: this is blocking, and - # ``load_to_graph`` runs it inside the connect/refresh streams. - response = batch_completion( + # Bounded and retried like every other provider call: one LLM_TIMEOUT + # budget for the batch, transient failures retried with what remains. + # (The raw litellm knobs are a trap here: it treats ``num_retries`` as + # overriding ``max_retries``, so passing the pair retried nothing.) + response = run_batch_completion( + batch_messages, model=Config.COMPLETION_MODEL, - messages=batch_messages, temperature=0, max_tokens=50, - timeout=Config.LLM_TIMEOUT, - max_retries=Config.LLM_MAX_RETRIES, - num_retries=0, + label="table-descriptions", ) for offset, batch_response in enumerate(response): diff --git a/tests/test_deadline_guard.py b/tests/test_deadline_guard.py index 24841ef9..350f87ff 100644 --- a/tests/test_deadline_guard.py +++ b/tests/test_deadline_guard.py @@ -4,9 +4,16 @@ ``tcp_user_timeout`` bounds unacknowledged outbound data, and keepalives only detect a dead TCP peer. A stalled backend or proxy satisfies all three while the client stays blocked in a read, holding a worker thread that cancellation cannot -reclaim. The guard cancels and then closes, so the read raises. +reclaim. The guard cancels and then shuts the socket down, so the read raises. + +The enforcement itself must not be blockable by the stall it breaks: psycopg2 +holds a connection lock during a blocking read, so ``close()`` from another +thread deadlocks, and ``cancel()`` (PQcancel) can hang against a black-holed +server. The tests below reproduce both against a real blocked socket read. """ +import os +import socket import threading import time @@ -76,3 +83,135 @@ def test_guard_is_a_noop_without_a_deadline(seconds): pass assert not conn.cancelled.is_set() assert not conn.closed.is_set() + + +class _SocketConn: + """Fake driver connection over a real socket, with psycopg2's locking. + + ``close()`` takes the same lock the worker holds while blocked in ``recv``, + which is exactly why the guard must not go through ``close()`` to break a + stall — this fake deadlocks if it tries. + """ + + def __init__(self, sock, lock, cancel_hangs=False): + self._sock = sock + self._lock = lock + self._cancel_hangs = cancel_hangs + self._cancel_release = threading.Event() + self.cancelled = threading.Event() + + def fileno(self): + return self._sock.fileno() + + def cancel(self): + self.cancelled.set() + if self._cancel_hangs: + # PQcancel against a black-holed server: a blocking connect with + # no timeout of ours to bound it. + self._cancel_release.wait(timeout=5) + + def close(self): + with self._lock: # would deadlock while the worker is blocked + self._sock.close() + + def release_cancel(self): + self._cancel_release.set() + + +def _blocked_reader(sock, lock, released): + """A worker blocked in a socket read, holding the driver lock.""" + + def worker(): + with lock: + try: + sock.recv(1) # the peer never answers + except OSError: + pass + released.set() + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + # Don't start the deadline until the worker actually holds the lock. + while not lock.locked(): + time.sleep(0.001) + return thread + + +@pytest.mark.unit +def test_guard_releases_a_read_that_close_would_deadlock_on(monkeypatch): + """The reported freeze: the worker blocked in recv holds the driver lock, + so a close() from the timer thread would join the deadlock. The socket + shutdown must free the worker anyway.""" + monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) + left, right = socket.socketpair() + lock = threading.Lock() + conn = _SocketConn(left, lock) + released = threading.Event() + + try: + _blocked_reader(left, lock, released) + with deadline.deadline_guard(conn, 0.1, "probe"): + assert released.wait(timeout=5), "the blocked read was never released" + assert conn.cancelled.is_set(), "cancel should still be attempted first" + finally: + right.close() + with lock: + left.close() + + +@pytest.mark.unit +def test_guard_does_not_wait_for_a_hanging_cancel(monkeypatch): + """PQcancel can hang against a black-holed server; the escalation to the + socket shutdown must proceed on its own clock regardless.""" + monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) + left, right = socket.socketpair() + lock = threading.Lock() + conn = _SocketConn(left, lock, cancel_hangs=True) + released = threading.Event() + + try: + _blocked_reader(left, lock, released) + started = time.monotonic() + with deadline.deadline_guard(conn, 0.1, "probe"): + assert released.wait(timeout=5), "shutdown waited on the hanging cancel" + assert time.monotonic() - started < 3, "release took far longer than deadline+grace" + finally: + conn.release_cancel() + right.close() + with lock: + left.close() + + +@pytest.mark.unit +def test_guard_shutdown_leaves_the_drivers_descriptor_valid(monkeypatch): + """shutdown(2) runs on a duplicated descriptor: the socket dies, but the + driver's own fd must stay valid for its cleanup path to close.""" + monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.1) + left, right = socket.socketpair() + lock = threading.Lock() + conn = _SocketConn(left, lock) + released = threading.Event() + + try: + _blocked_reader(left, lock, released) + with deadline.deadline_guard(conn, 0.1, "probe"): + assert released.wait(timeout=5) + os.fstat(left.fileno()) # raises if the guard closed the driver's fd + finally: + right.close() + with lock: + left.close() + + +@pytest.mark.unit +def test_guard_falls_back_to_close_without_a_socket(monkeypatch): + """A connection that exposes no usable descriptor still gets closed.""" + monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.1) + + class _NoFdConn(_Conn): + def fileno(self): + raise RuntimeError("connection already closed") + + conn = _NoFdConn() + with deadline.deadline_guard(conn, 0.1, "probe"): + assert conn.closed.wait(timeout=2), "close fallback did not run" diff --git a/tests/test_llm_retry_policy.py b/tests/test_llm_retry_policy.py new file mode 100644 index 00000000..cc045df0 --- /dev/null +++ b/tests/test_llm_retry_policy.py @@ -0,0 +1,200 @@ +"""Retries must be spent on failures a replay can change. + +A 401 does not become a valid key by asking twice, and a transient 500 on a +batch item should not silently degrade a table description to its name. One +policy decides both: ``_retryable`` judges the failure, ``run_completion`` and +``run_batch_completion`` apply the verdict against the remaining budget. +""" + +import types + +import pytest + +import api.agents.utils as agent_utils + +pytestmark = pytest.mark.unit + + +class _ProviderError(Exception): + """A provider failure carrying an HTTP verdict, litellm-style.""" + + def __init__(self, status_code=None, retry_after=None): + super().__init__(f"status={status_code}") + if status_code is not None: + self.status_code = status_code + if retry_after is not None: + self.response = types.SimpleNamespace( + headers={"retry-after": retry_after} + ) + + +def _ok(content="ok"): + message = types.SimpleNamespace(content=content) + choice = types.SimpleNamespace(message=message) + return types.SimpleNamespace(choices=[choice]) + + +@pytest.fixture(autouse=True) +def _fast_backoff(monkeypatch): + monkeypatch.setattr(agent_utils, "_RETRY_BACKOFF_SECONDS", 0.001) + + +class TestRetryVerdicts: + """The classification itself, one failure class at a time.""" + + @pytest.mark.parametrize("status_code", [400, 401, 403, 404, 422]) + def test_the_request_itself_being_bad_is_permanent(self, status_code): + assert agent_utils._retryable(_ProviderError(status_code)) is False + + @pytest.mark.parametrize("status_code", [408, 429, 500, 502, 503, 504]) + def test_the_service_having_a_moment_is_transient(self, status_code): + assert agent_utils._retryable(_ProviderError(status_code)) is True + + def test_a_failure_with_no_verdict_is_transient(self): + """Timeouts and dropped connections carry no status code.""" + assert agent_utils._retryable(RuntimeError("connection reset")) is True + + def test_retry_after_wins_over_backoff(self): + delay = agent_utils._retry_delay( + _ProviderError(429, retry_after="0.25"), attempt=1 + ) + assert delay == 0.25 + + def test_retry_after_is_capped(self): + """A provider asking for an hour cannot eat the whole budget.""" + delay = agent_utils._retry_delay( + _ProviderError(429, retry_after="3600"), attempt=1 + ) + assert delay <= agent_utils._RETRY_BACKOFF_CAP_SECONDS + + def test_garbage_retry_after_falls_back_to_backoff(self): + delay = agent_utils._retry_delay( + _ProviderError(429, retry_after="tomorrow"), attempt=1 + ) + assert delay == agent_utils._RETRY_BACKOFF_SECONDS + + +class TestRunCompletion: + """The single-call loop applies the verdicts.""" + + def test_a_401_makes_exactly_one_request(self, monkeypatch): + monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", 3, + raising=False) + calls = [] + + def unauthorized(**kwargs): + calls.append(kwargs) + raise _ProviderError(401) + + monkeypatch.setattr(agent_utils, "completion", unauthorized) + + with pytest.raises(_ProviderError): + agent_utils.run_completion([{"role": "user", "content": "hi"}], + label="probe") + + assert len(calls) == 1, "a permanent failure was replayed" + + def test_a_transient_failure_is_retried(self, monkeypatch): + monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", 1, + raising=False) + calls = [] + + def flaky(**kwargs): + calls.append(kwargs) + if len(calls) == 1: + raise _ProviderError(500) + return _ok("recovered") + + monkeypatch.setattr(agent_utils, "completion", flaky) + + result = agent_utils.run_completion( + [{"role": "user", "content": "hi"}], label="probe") + + assert result == "recovered" + assert len(calls) == 2 + + +class TestRunBatchCompletion: + """The batch loop retries only the slots worth retrying.""" + + def test_only_transient_failures_are_replayed(self, monkeypatch): + monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", 2, + raising=False) + permanent = _ProviderError(401) + transient = _ProviderError(500) + batches = [] + + def fake_batch(**kwargs): + batches.append(kwargs["messages"]) + if len(batches) == 1: + return [_ok("first"), transient, permanent] + return [_ok("second")] # only the transient slot came back + + monkeypatch.setattr(agent_utils, "batch_completion", fake_batch) + + messages = [[{"role": "user", "content": f"table {i}"}] + for i in range(3)] + results = agent_utils.run_batch_completion(messages, model="m", + label="probe") + + assert len(batches) == 2 + assert batches[1] == [messages[1]], "retry was not scoped to the transient slot" + assert results[0].choices[0].message.content == "first" + assert results[1].choices[0].message.content == "second" + assert results[2] is permanent, "the permanent failure must be kept, not retried" + + def test_a_clean_batch_makes_one_call(self, monkeypatch): + calls = [] + + def fake_batch(**kwargs): + calls.append(kwargs) + return [_ok(), _ok()] + + monkeypatch.setattr(agent_utils, "batch_completion", fake_batch) + + results = agent_utils.run_batch_completion( + [[{"role": "user", "content": "a"}], + [{"role": "user", "content": "b"}]], + model="m", label="probe") + + assert len(calls) == 1 + assert len(results) == 2 + # Library retry knobs stay off; the budget bounds the call. + assert calls[0]["max_retries"] == 0 + assert calls[0]["num_retries"] == 0 + assert calls[0]["timeout"] <= agent_utils.Config.LLM_TIMEOUT + + def test_attempts_run_against_the_remaining_budget(self, monkeypatch): + """A later attempt gets less time, never a fresh allocation.""" + monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", 2, + raising=False) + timeouts = [] + + def failing_batch(**kwargs): + timeouts.append(kwargs["timeout"]) + return [_ProviderError(500)] + + monkeypatch.setattr(agent_utils, "batch_completion", failing_batch) + + results = agent_utils.run_batch_completion( + [[{"role": "user", "content": "a"}]], model="m", label="probe") + + assert len(timeouts) == 3 + assert timeouts == sorted(timeouts, reverse=True) + assert isinstance(results[0], _ProviderError) + + def test_a_slot_the_library_never_answered_is_a_failure(self, monkeypatch): + """A short batch response must not surface as a None 'success'.""" + + def short_batch(**kwargs): + return [_ok("only one")] # two were asked for + + monkeypatch.setattr(agent_utils, "batch_completion", short_batch) + + results = agent_utils.run_batch_completion( + [[{"role": "user", "content": "a"}], + [{"role": "user", "content": "b"}]], + model="m", label="probe") + + assert results[0].choices[0].message.content == "only one" + assert isinstance(results[1], Exception), "an unanswered slot leaked as None" diff --git a/tests/test_settings_route.py b/tests/test_settings_route.py index 9c23a654..2095f042 100644 --- a/tests/test_settings_route.py +++ b/tests/test_settings_route.py @@ -121,9 +121,9 @@ async def test_valid_key_returns_success(self, mock_completion, mock_request): messages=[{"role": "user", "content": "test"}], max_tokens=1, api_key="sk-validkey123456", - timeout=Config.LLM_TIMEOUT, - max_retries=Config.LLM_MAX_RETRIES, - num_retries=0, + # One bounded attempt via the shared helper: a 401 is the + # probe's answer, so library retries stay off. + **Config.llm_call_bounds(), ) @pytest.mark.asyncio @@ -177,9 +177,9 @@ async def test_gemini_vendor_accepted(self, mock_completion, mock_request): messages=[{"role": "user", "content": "test"}], max_tokens=1, api_key="AIzaSyTest123456", - timeout=Config.LLM_TIMEOUT, - max_retries=Config.LLM_MAX_RETRIES, - num_retries=0, + # One bounded attempt via the shared helper: a 401 is the + # probe's answer, so library retries stay off. + **Config.llm_call_bounds(), ) @pytest.mark.asyncio diff --git a/tests/test_timeout_validation.py b/tests/test_timeout_validation.py index 95665d56..e94c25d3 100644 --- a/tests/test_timeout_validation.py +++ b/tests/test_timeout_validation.py @@ -131,6 +131,9 @@ def test_run_completion_makes_exactly_the_budgeted_attempts(monkeypatch, retries monkeypatch.setattr(agent_utils.Config, "LLM_TIMEOUT", 5.0, raising=False) monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", retries, raising=False) + # Backoff between retries is real behaviour but slow in a test. + monkeypatch.setattr(agent_utils, "_RETRY_BACKOFF_SECONDS", 0.001) + calls = [] def failing_completion(**kwargs): @@ -157,6 +160,8 @@ def test_run_completion_stops_retrying_when_the_budget_is_spent(monkeypatch): monkeypatch.setattr(agent_utils.Config, "LLM_TIMEOUT", 0.3, raising=False) monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", 5, raising=False) + monkeypatch.setattr(agent_utils, "_RETRY_BACKOFF_SECONDS", 0.001) + calls = [] def slow_failing_completion(**_kwargs): From 676d5122baf027be780a57d0b646c76429a53c23 Mon Sep 17 00:00:00 2001 From: Gal Shubeli Date: Sun, 23 Aug 2026 18:09:12 +0300 Subject: [PATCH 25/25] fix: drop in-process PQcancel; honour the provider's Retry-After MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the two findings from @Naseem77's latest review, then stops here — see the PR comment on scope. **1. The deadline guard no longer calls `conn.cancel()`.** PQcancel opens its own connection to the server, so against a black-holed host it hangs, and psycopg2 does not release the GIL around it — a hanging cancel stalls every Python thread in the process, including the event loop and so every stream keepalive. Running it on a separate daemon timer was no protection: a thread that cannot acquire the GIL cannot run, which is why his probe saw socket shutdown at 5.1s for a 1s deadline with a 0.5s grace. The escalation was already a `shutdown(2)` on a duplicated descriptor, and that is what actually releases the blocked read, so the cancel was removed rather than replaced. There is now one timer firing at the deadline and no grace period, so the deadline is the whole clock. Skipping the cancel costs little: the backend aborts the query itself once it notices the client socket is gone. This is a net deletion, and it removes a failure mode worse than the gap it was closing. **2. `Retry-After` is read from litellm's headers and honoured in full.** litellm keeps the provider's original headers as `litellm_response_headers`; `exc.response.headers` is its own reconstructed response and carries none, so every real 429 fell through to the exponential guess. Measured against a live `429 Retry-After: 10` with a 15s budget: before: 2 requests, 0.52s apart (the 0.5s fallback) after: 2 requests, 10.02s apart, 10.19s total The delay is no longer capped at 8s either. Capping produced a retry certain to be refused again, whereas the caller already declines a delay that outlives the budget — so a long `Retry-After` now means no retry rather than a premature one. Verified: same 429 with a 5s budget makes exactly 1 request. HTTP-date form is accepted alongside seconds. An existing test asserted the capping behaviour and is updated to the new contract. 362 unit + 14 SDK tests pass; pylint 10.00/10; make lint clean. Refs: research#86 Co-Authored-By: Claude Opus 5 (1M context) --- api/agents/utils.py | 58 +++++++++++++++++---- api/loaders/deadline.py | 86 ++++++++++++-------------------- tests/test_deadline_guard.py | 57 ++++++++++++--------- tests/test_llm_retry_policy.py | 13 +++-- tests/test_timeout_validation.py | 64 ++++++++++++++++++++++++ 5 files changed, 188 insertions(+), 90 deletions(-) diff --git a/api/agents/utils.py b/api/agents/utils.py index 62eca64c..1d1d13b7 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -3,6 +3,8 @@ import json import logging import time +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime from typing import Any, Dict, List from litellm import batch_completion, completion @@ -46,28 +48,66 @@ def _retryable(exc: Exception) -> bool: return status_code in (408, 429) or status_code >= 500 -def _retry_after_seconds(exc: Exception) -> float | None: - """The provider's Retry-After, if the failure carried one.""" - headers = getattr(getattr(exc, "response", None), "headers", None) +def _retry_after_header(exc: Exception): + """The raw ``Retry-After`` value from wherever the client stashed it.""" + for source in ("litellm_response_headers", "response_headers"): + headers = getattr(exc, source, None) + if headers is not None: + break + else: + headers = getattr(getattr(exc, "response", None), "headers", None) if headers is None: return None try: - value = headers.get("retry-after") - return float(value) if value is not None else None + return headers.get("retry-after") + except (AttributeError, TypeError): + return None + + +def _http_date_delay(value) -> float | None: + """Seconds until an HTTP-date ``Retry-After``, the spec's other form.""" + try: + when = parsedate_to_datetime(str(value)) except (TypeError, ValueError): return None + if when is None: + return None + if when.tzinfo is None: + when = when.replace(tzinfo=timezone.utc) + return max(0.0, (when - datetime.now(timezone.utc)).total_seconds()) + + +def _retry_after_seconds(exc: Exception) -> float | None: + """The provider's ``Retry-After``, if the failure carried one. + + litellm keeps the provider's original headers on the exception as + ``litellm_response_headers``; ``exc.response.headers`` is litellm's own + reconstructed response and does not carry them, so reading that returned + ``None`` for every real 429 and the backoff fell back to its own guess. + + Both header spellings are accepted: a delay in seconds, or an HTTP date. + """ + value = _retry_after_header(exc) + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return _http_date_delay(value) def _retry_delay(exc: Exception, attempt: int) -> float: """Seconds to wait before the next attempt. - The provider's own Retry-After wins when present (a 429 tells us exactly - when trying again stops being rude); otherwise a small exponential - backoff, so a struggling service is not hammered at full speed. + The provider's own ``Retry-After`` wins when present — a 429 states exactly + when trying again stops being rude — and is honoured in full rather than + truncated: the caller already refuses a delay that outlives the budget, so + capping it here only produced a retry that was certain to be rejected + again. The exponential fallback stays capped, since it is a guess. """ retry_after = _retry_after_seconds(exc) if retry_after is not None and retry_after >= 0: - return min(retry_after, _RETRY_BACKOFF_CAP_SECONDS) + return retry_after return min(_RETRY_BACKOFF_SECONDS * (2 ** (attempt - 1)), _RETRY_BACKOFF_CAP_SECONDS) diff --git a/api/loaders/deadline.py b/api/loaders/deadline.py index e47bc192..0dffee46 100644 --- a/api/loaders/deadline.py +++ b/api/loaders/deadline.py @@ -7,26 +7,31 @@ alive without answering satisfies all three while the client stays blocked in a read — holding a worker thread that cancellation cannot reclaim. -This closes that gap from the outside: a timer asks the server to cancel the -in-flight query and, failing that, shuts the connection's socket down at the OS -level, which makes the blocked read raise in the worker so the thread is -released. - -The escalation must not itself be blockable by the stall it exists to break: - -* ``conn.cancel()`` (PQcancel) opens its *own* connection to the server, so - against a black-holed host the cancel itself can hang. It therefore runs on - a dedicated daemon timer thread that nothing waits for — a hanging cancel - costs one parked thread until the kernel gives up, never the escalation. +This closes that gap from the outside: a timer shuts the connection's socket +down at the OS level, which makes the blocked read raise in the worker so the +thread is released. + +The escalation must not itself be blockable by the stall it exists to break, +which rules out both cooperative options: + +* ``conn.cancel()`` (PQcancel) is deliberately *not* used. It opens its own + connection to the server, so against a black-holed host the cancel itself + hangs — and psycopg2 does not release the GIL around it, so a hanging cancel + stalls every Python thread in the process, including the event loop and so + every stream keepalive. Running it on a separate daemon thread does not help: + a thread that cannot acquire the GIL cannot run. Skipping the cancel costs + little, since the backend aborts the query itself once it notices the + client's socket has gone. * ``conn.close()`` is not usable from another thread while a query is running: psycopg2 serialises connection access with an internal lock that the worker blocked in ``recv`` is holding, so a close from the timer thread would join - the deadlock instead of breaking it. The escalation instead calls - ``shutdown(2)`` on the underlying socket via a duplicated file descriptor — - a plain syscall that needs no driver lock and no network round trip. The - kernel fails the pending read immediately, the driver raises in the worker, - and the connection is then closed by the worker's own cleanup path, which - holds the lock legitimately. + the deadlock instead of breaking it. + +What is left is ``shutdown(2)`` on the underlying socket via a duplicated file +descriptor — a plain syscall that needs no driver lock, no network round trip +and no GIL-holding driver call. The kernel fails the pending read immediately, +the driver raises in the worker, and the connection is then closed by the +worker's own cleanup path, which holds the lock legitimately. """ import contextlib @@ -35,29 +40,6 @@ import socket import threading -# How long to wait for a cooperative cancel before shutting the socket down. -_CANCEL_GRACE_SECONDS = 5.0 - - -def _cancel(conn, label: str) -> None: - """Ask the server to abort the running statement, if the driver can. - - Best-effort only: this may hang (see the module docstring), and the - escalation to ``_shutdown`` does not depend on it returning. - """ - cancel = getattr(conn, "cancel", None) - if cancel is None: - return - logging.warning("%s exceeded its deadline; cancelling the query", label) - try: - cancel() - except Exception as exc: # pylint: disable=broad-exception-caught - # PQcancel opens its own connection to the server, so it can fail when - # the server is unreachable. The socket shutdown is the fallback. - logging.warning("%s cancel failed (%s); the socket shutdown will follow", - label, type(exc).__name__) - - def _connection_fd(conn): """The connection's socket descriptor, or ``None`` if it has none.""" fileno = getattr(conn, "fileno", None) @@ -120,12 +102,12 @@ def _close(conn, label: str) -> None: @contextlib.contextmanager def deadline_guard(conn, seconds: float, label: str = "database call"): - """Cancel, then shut down, *conn* if the body outlives *seconds*. + """Shut *conn* down if the body outlives *seconds*. The guard is what makes the configured deadline real for a connection whose - peer has stopped answering but has not dropped the socket. Both escalation - steps run on their own daemon timer threads and neither waits for the - other, so a step that itself hangs cannot postpone the next one. + peer has stopped answering but has not dropped the socket. The escalation + runs on a daemon timer and performs only a syscall, so nothing on the path + can be blocked by the stall it exists to break. """ if not seconds or seconds <= 0: yield @@ -142,19 +124,15 @@ def deadline_guard(conn, seconds: float, label: str = "database call"): # preferred whenever it exists. escalate, escalate_args = _close, (conn, label) - cancel_timer = threading.Timer(seconds, _cancel, args=(conn, label)) - shutdown_timer = threading.Timer( - seconds + _CANCEL_GRACE_SECONDS, escalate, args=escalate_args - ) - cancel_timer.daemon = True - shutdown_timer.daemon = True - cancel_timer.start() - shutdown_timer.start() + # One timer, firing at the deadline: there is no cooperative step to wait + # for, so there is no grace period either. + escalation = threading.Timer(seconds, escalate, args=escalate_args) + escalation.daemon = True + escalation.start() try: yield finally: - cancel_timer.cancel() - shutdown_timer.cancel() + escalation.cancel() if dup is not None: # Closing the duplicate also disarms a shutdown callback that # already started: the socket object refuses use after close, so diff --git a/tests/test_deadline_guard.py b/tests/test_deadline_guard.py index 350f87ff..30764464 100644 --- a/tests/test_deadline_guard.py +++ b/tests/test_deadline_guard.py @@ -42,29 +42,36 @@ def close(self): @pytest.mark.unit -def test_guard_cancels_then_closes_a_stalled_call(monkeypatch): - monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) +def test_guard_escalates_a_stalled_call(): conn = _Conn() with deadline.deadline_guard(conn, 0.1, "probe"): # Stands in for a read that never returns. - assert conn.cancelled.wait(timeout=2), "deadline did not cancel the query" - assert conn.closed.wait(timeout=2), "deadline did not close the connection" + assert conn.closed.wait(timeout=2), "deadline did not release the call" @pytest.mark.unit -def test_guard_closes_even_when_cancel_fails(monkeypatch): - """PQcancel opens its own connection, so it can fail on an unreachable server.""" - monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) - conn = _Conn(cancel_raises=True) +def test_guard_never_calls_the_drivers_cancel(): + """PQcancel is deliberately unused. + + It opens its own connection to the server, so against a black-holed host it + hangs — and psycopg2 does not release the GIL around it, which stalls every + Python thread in the process, the event loop included. A separate timer + thread is no protection: a thread that cannot take the GIL cannot run. + """ + conn = _Conn() with deadline.deadline_guard(conn, 0.1, "probe"): - assert conn.closed.wait(timeout=2), "close fallback did not run" + assert conn.closed.wait(timeout=2), "deadline did not release the call" + + assert not conn.cancelled.is_set(), ( + "cancel was invoked in-process; a hanging PQcancel would freeze every " + "thread, including the streams this deadline exists to protect" + ) @pytest.mark.unit -def test_guard_leaves_a_prompt_call_alone(monkeypatch): - monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.1) +def test_guard_leaves_a_prompt_call_alone(): conn = _Conn() with deadline.deadline_guard(conn, 5.0, "probe"): @@ -138,11 +145,10 @@ def worker(): @pytest.mark.unit -def test_guard_releases_a_read_that_close_would_deadlock_on(monkeypatch): +def test_guard_releases_a_read_that_close_would_deadlock_on(): """The reported freeze: the worker blocked in recv holds the driver lock, so a close() from the timer thread would join the deadlock. The socket shutdown must free the worker anyway.""" - monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) left, right = socket.socketpair() lock = threading.Lock() conn = _SocketConn(left, lock) @@ -152,7 +158,7 @@ def test_guard_releases_a_read_that_close_would_deadlock_on(monkeypatch): _blocked_reader(left, lock, released) with deadline.deadline_guard(conn, 0.1, "probe"): assert released.wait(timeout=5), "the blocked read was never released" - assert conn.cancelled.is_set(), "cancel should still be attempted first" + assert not conn.cancelled.is_set(), "cancel must not be attempted" finally: right.close() with lock: @@ -160,10 +166,13 @@ def test_guard_releases_a_read_that_close_would_deadlock_on(monkeypatch): @pytest.mark.unit -def test_guard_does_not_wait_for_a_hanging_cancel(monkeypatch): - """PQcancel can hang against a black-holed server; the escalation to the - socket shutdown must proceed on its own clock regardless.""" - monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.15) +def test_guard_escalates_on_the_deadline_alone(): + """The deadline is the whole clock: nothing cooperative precedes it. + + An earlier version cancelled first and shut the socket down only after a + grace period, which made the real release time depend on how long PQcancel + took — and a GIL-holding cancel could postpone it indefinitely. + """ left, right = socket.socketpair() lock = threading.Lock() conn = _SocketConn(left, lock, cancel_hangs=True) @@ -173,8 +182,10 @@ def test_guard_does_not_wait_for_a_hanging_cancel(monkeypatch): _blocked_reader(left, lock, released) started = time.monotonic() with deadline.deadline_guard(conn, 0.1, "probe"): - assert released.wait(timeout=5), "shutdown waited on the hanging cancel" - assert time.monotonic() - started < 3, "release took far longer than deadline+grace" + assert released.wait(timeout=5), "the blocked read was never released" + elapsed = time.monotonic() - started + assert elapsed < 1.0, f"release took {elapsed:.2f}s for a 0.1s deadline" + assert not conn.cancelled.is_set(), "cancel must not be attempted" finally: conn.release_cancel() right.close() @@ -183,10 +194,9 @@ def test_guard_does_not_wait_for_a_hanging_cancel(monkeypatch): @pytest.mark.unit -def test_guard_shutdown_leaves_the_drivers_descriptor_valid(monkeypatch): +def test_guard_shutdown_leaves_the_drivers_descriptor_valid(): """shutdown(2) runs on a duplicated descriptor: the socket dies, but the driver's own fd must stay valid for its cleanup path to close.""" - monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.1) left, right = socket.socketpair() lock = threading.Lock() conn = _SocketConn(left, lock) @@ -204,9 +214,8 @@ def test_guard_shutdown_leaves_the_drivers_descriptor_valid(monkeypatch): @pytest.mark.unit -def test_guard_falls_back_to_close_without_a_socket(monkeypatch): +def test_guard_falls_back_to_close_without_a_socket(): """A connection that exposes no usable descriptor still gets closed.""" - monkeypatch.setattr(deadline, "_CANCEL_GRACE_SECONDS", 0.1) class _NoFdConn(_Conn): def fileno(self): diff --git a/tests/test_llm_retry_policy.py b/tests/test_llm_retry_policy.py index cc045df0..a6e61972 100644 --- a/tests/test_llm_retry_policy.py +++ b/tests/test_llm_retry_policy.py @@ -60,12 +60,19 @@ def test_retry_after_wins_over_backoff(self): ) assert delay == 0.25 - def test_retry_after_is_capped(self): - """A provider asking for an hour cannot eat the whole budget.""" + def test_retry_after_is_honoured_in_full(self): + """A long delay is reported as-is; the caller decides whether it fits. + + Truncating it here produced a retry that was certain to be refused + again — and the budget check already declines a delay it cannot afford, + so an hour-long Retry-After simply means no retry rather than a capped + one. + """ delay = agent_utils._retry_delay( _ProviderError(429, retry_after="3600"), attempt=1 ) - assert delay <= agent_utils._RETRY_BACKOFF_CAP_SECONDS + assert delay == 3600.0 + assert delay > agent_utils._RETRY_BACKOFF_CAP_SECONDS def test_garbage_retry_after_falls_back_to_backoff(self): delay = agent_utils._retry_delay( diff --git a/tests/test_timeout_validation.py b/tests/test_timeout_validation.py index e94c25d3..5d6e1bf2 100644 --- a/tests/test_timeout_validation.py +++ b/tests/test_timeout_validation.py @@ -7,6 +7,7 @@ import importlib import time +import types import pytest @@ -184,3 +185,66 @@ def test_library_retry_knobs_are_disabled(): bounds = agent_utils.Config.llm_call_bounds(timeout=7) assert bounds == {"timeout": 7, "max_retries": 0, "num_retries": 0} + + +def _rate_limited(retry_after, source="litellm_response_headers"): + """A 429-style exception carrying its delay where litellm puts it.""" + exc = RuntimeError("rate limited") + setattr(exc, source, {"retry-after": retry_after}) + # litellm's own reconstructed response does NOT carry provider headers; + # reading it is what silently lost every Retry-After. + exc.response = types.SimpleNamespace(headers={}) + return exc + + +@pytest.mark.unit +def test_retry_after_is_read_from_litellms_headers(): + """The provider's delay lives on litellm_response_headers, not response.""" + import api.agents.utils as agent_utils + + assert agent_utils._retry_after_seconds(_rate_limited("10")) == 10.0 + # Nothing to honour when no header is present anywhere. + assert agent_utils._retry_after_seconds(RuntimeError("boom")) is None + + +@pytest.mark.unit +def test_retry_after_accepts_an_http_date(): + import api.agents.utils as agent_utils + from email.utils import format_datetime + from datetime import datetime, timedelta, timezone + + when = datetime.now(timezone.utc) + timedelta(seconds=30) + delay = agent_utils._retry_after_seconds(_rate_limited(format_datetime(when))) + assert delay is not None and 25 <= delay <= 31 + + +@pytest.mark.unit +def test_retry_after_is_honoured_in_full_not_capped(): + """Truncating the delay produced a retry certain to be refused again.""" + import api.agents.utils as agent_utils + + delay = agent_utils._retry_delay(_rate_limited("30"), attempt=1) + assert delay == 30.0 + assert delay > agent_utils._RETRY_BACKOFF_CAP_SECONDS + + +@pytest.mark.unit +def test_retry_is_skipped_when_the_delay_will_not_fit(monkeypatch): + """A delay longer than the remaining budget means no second attempt.""" + import api.agents.utils as agent_utils + + monkeypatch.setattr(agent_utils.Config, "LLM_TIMEOUT", 2.0, raising=False) + monkeypatch.setattr(agent_utils.Config, "LLM_MAX_RETRIES", 3, raising=False) + + calls = [] + + def rate_limited_completion(**_kwargs): + calls.append(1) + raise _rate_limited("30") + + monkeypatch.setattr(agent_utils, "completion", rate_limited_completion) + + with pytest.raises(RuntimeError, match="rate limited"): + agent_utils.run_completion([{"role": "user", "content": "hi"}], label="probe") + + assert len(calls) == 1, "slept-or-retried past a delay that could not fit"