diff --git a/.env.example b/.env.example index eeecd753..8f5f1826 100644 --- a/.env.example +++ b/.env.example @@ -75,6 +75,34 @@ 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 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). +# 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 +# +# 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 +# +# 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 @@ -100,7 +128,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/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/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..6f46b57b 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -1,30 +1,115 @@ """Utility functions for agents.""" import json +import logging +import time from typing import Any, Dict, List from litellm import completion from api.config import Config -def run_completion(messages: List[Dict[str, str]], custom_model: str = None, - custom_api_key: str = None, **kwargs) -> str: +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: """Run an LLM completion with optional custom model/key overrides. + 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. """ - completion_args = { + base_args = { "model": custom_model if custom_model else Config.COMPLETION_MODEL, "messages": messages, "top_p": 1, - **kwargs, } - if custom_api_key: - completion_args["api_key"] = custom_api_key + base_args["api_key"] = custom_api_key - result = completion(**completion_args) - return result.choices[0].message.content + 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 + ) + + 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" + ) class BaseAgent: # pylint: disable=too-few-public-methods diff --git a/api/config.py b/api/config.py index dce8a34c..f4ab056d 100644 --- a/api/config.py +++ b/api/config.py @@ -4,7 +4,9 @@ """ import os +import time import logging +import math import dataclasses from typing import Union @@ -40,10 +42,21 @@ 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 Config.llm_call_bounds() + 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 +64,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 +86,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 @@ -77,8 +109,38 @@ 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 + # ``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 " + "timeout disables the safeguard it exists to provide" + ) + return value + + @dataclasses.dataclass -class Config: +class Config: # pylint: disable=too-many-instance-attributes """ Configuration class for the text2sql module. """ @@ -133,6 +195,80 @@ 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 = _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 = _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 + # 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. + # Zero is valid here (it means "no retry", a strict ceiling); negative is + # not. + # pylint: disable-next=invalid-name + LLM_MAX_RETRIES: int = max(0, int(os.getenv("LLM_MAX_RETRIES", "1"))) + + @classmethod + 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. + """ + return { + "timeout": cls.LLM_TIMEOUT if timeout is None else timeout, + "max_retries": 0, + "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. + # pylint: disable-next=invalid-name + DB_CONNECT_TIMEOUT: int = _positive_env("DB_CONNECT_TIMEOUT", "10") + # 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/core/text2sql.py b/api/core/text2sql.py index 90cd9cdc..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,18 +386,39 @@ 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( 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 +438,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, ) @@ -513,7 +525,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", @@ -536,7 +553,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 +614,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, @@ -742,7 +764,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: @@ -755,7 +780,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/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/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)] 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/graph_loader.py b/api/loaders/graph_loader.py index b4f3fca5..ce08d3f6 100644 --- a/api/loaders/graph_loader.py +++ b/api/loaders/graph_loader.py @@ -1,11 +1,12 @@ """Graph loader module for loading data into graph databases.""" +import asyncio import json 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,10 +31,17 @@ 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() - - create_combined_description(entities) + # 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() + + # 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 @@ -56,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 { @@ -70,7 +80,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 +119,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 +133,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/introspection.py b/api/loaders/introspection.py new file mode 100644 index 00000000..a42c5b55 --- /dev/null +++ b/api/loaders/introspection.py @@ -0,0 +1,60 @@ +"""Bounded execution for schema introspection. + +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 + +_EXECUTOR: ThreadPoolExecutor | None = None +_EXECUTOR_LOCK = threading.Lock() + + +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* on the bounded introspection pool. + + 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. + """ + pool = _executor() + queued = getattr(pool, "_work_queue", None) + if queued is not None and queued.qsize(): + logging.info( + "schema introspection queued: all %d workers busy", + Config.DB_SCHEMA_CONCURRENCY, + ) + + 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/mysql_loader.py b/api/loaders/mysql_loader.py index 2e8b40fa..fdeaeb19 100644 --- a/api/loaders/mysql_loader.py +++ b/api/loaders/mysql_loader.py @@ -11,8 +11,10 @@ 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 +from api.loaders.introspection import run_introspection class MySQLQueryError(Exception): @@ -151,6 +153,42 @@ 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: + # 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, + 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) + 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, @@ -171,25 +209,12 @@ async def load( # pylint: disable=arguments-differ try: # Parse connection URL conn_params = MySQLLoader._parse_mysql_url(connection_url) - - # Connect to MySQL database - conn = pymysql.connect(**conn_params) - cursor = conn.cursor(DictCursor) - - # Get database name db_name = conn_params['database'] - # Get all table information yield True, "Extracting table information..." - entities = MySQLLoader.extract_tables_info(cursor, db_name) - - # Get all relationship information - yield True, "Extracting relationship information..." - relationships = MySQLLoader.extract_relationships(cursor, db_name) - - # Close database connection - cursor.close() - conn.close() + entities, relationships = await run_introspection( + MySQLLoader._introspect_schema, conn_params, db_name + ) # Load data into graph yield True, "Loading data into graph..." @@ -512,6 +537,14 @@ 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. ``_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) cursor = conn.cursor(DictCursor) diff --git a/api/loaders/postgres_loader.py b/api/loaders/postgres_loader.py index c5dff6fa..8dc6529f 100644 --- a/api/loaders/postgres_loader.py +++ b/api/loaders/postgres_loader.py @@ -11,8 +11,27 @@ 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.deadline import deadline_guard 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. +# ``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 + # 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*)" +) logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") @@ -139,6 +158,117 @@ 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 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*([+-]?)\s*" # optional sign + 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, + ) + if not match: + return None + + 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 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. + 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 or "ms").lower()] + 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. + + 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: + # 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. + # 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, + **PostgresLoader._connect_kwargs( + connection_url, Config.DB_SCHEMA_TIMEOUT + ), + ) + cursor = conn.cursor() + + # 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: + cursor.close() + if conn is not None: + conn.close() + @staticmethod async def load( # pylint: disable=arguments-differ prefix: str, @@ -158,40 +288,19 @@ 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) - # Connect to PostgreSQL database - conn = psycopg2.connect(connection_url) - 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)) - ) - # 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 = PostgresLoader.extract_tables_info(cursor, schema) - - yield True, "Extracting relationship information..." - # Get all relationship information - relationships = PostgresLoader.extract_relationships(cursor, schema) - - # Close database connection before graph loading - cursor.close() - cursor = None - conn.close() - conn = None + entities, relationships = await run_introspection( + PostgresLoader._introspect_schema, connection_url, schema + ) yield True, "Loading data into graph..." # Load data into graph @@ -207,11 +316,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]: @@ -536,6 +640,89 @@ 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]: + """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 + 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] = {} + + # 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 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 = statement_timeout_s * 1000 + 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 + + 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 + + # 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. + # 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 def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: """ @@ -550,12 +737,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) + 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 - 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/api/loaders/snowflake_loader.py b/api/loaders/snowflake_loader.py index 7685daa9..4fa44881 100644 --- a/api/loaders/snowflake_loader.py +++ b/api/loaders/snowflake_loader.py @@ -15,8 +15,10 @@ 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 +from api.loaders.introspection import run_introspection class SnowflakeQueryError(Exception): @@ -241,7 +243,50 @@ def _parse_snowflake_url(connection_url: str) -> Dict[str, Any]: # pylint: disa return conn_params @staticmethod - async def load(prefix: str, connection_url: str) -> AsyncGenerator[ + 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: + # 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 + # ``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, + } + 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, db=None) -> AsyncGenerator[ tuple[bool, str], None ]: """ @@ -257,33 +302,21 @@ async def load(prefix: str, connection_url: str) -> AsyncGenerator[ try: # Parse connection URL conn_params = SnowflakeLoader._parse_snowflake_url(connection_url) - - # Connect to Snowflake database - conn = 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 = 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) - - # Close database connection - cursor.close() - conn.close() + entities, relationships = await run_introspection( + SnowflakeLoader._introspect_schema, + conn_params, db_name, schema_name, + ) # 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.") @@ -578,7 +611,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. @@ -593,11 +628,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) @@ -612,7 +647,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: @@ -646,6 +681,20 @@ 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. + # ``_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 + # 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 + ) + conn_params["session_parameters"] = session_parameters + # Connect to Snowflake database conn = snowflake.connector.connect(**conn_params) cursor = conn.cursor(DictCursor) diff --git a/api/memory/graphiti_tool.py b/api/memory/graphiti_tool.py index 1a052c57..eb4e407c 100644 --- a/api/memory/graphiti_tool.py +++ b/api/memory/graphiti_tool.py @@ -27,7 +27,11 @@ from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF -from litellm import completion +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, @@ -253,7 +257,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 = [] @@ -261,14 +265,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 @@ -277,6 +281,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: @@ -352,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""" @@ -435,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 @@ -733,7 +740,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 = [] @@ -741,14 +748,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 } @@ -769,7 +774,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/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/routes/streaming.py b/api/routes/streaming.py new file mode 100644 index 00000000..cada6d37 --- /dev/null +++ b/api/routes/streaming.py @@ -0,0 +1,97 @@ +"""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. + + 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. + """ + # 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: + 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. + # CancelledError is a BaseException, so it is not caught here and + # propagates as cancellation should. + await queue.put(exc) + else: + await queue.put(finished) + + pump = asyncio.ensure_future(_pump()) + try: + while True: + 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: + return + if isinstance(item, Exception): + raise item + yield item + finally: + # Non-awaiting cleanup: safe even when cancellation is already pending. + pump.cancel() 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/app/src/components/chat/ChatInterface.tsx b/app/src/components/chat/ChatInterface.tsx index 3c9ed543..6bf9baa1 100644 --- a/app/src/components/chat/ChatInterface.tsx +++ b/app/src/components/chat/ChatInterface.tsx @@ -238,12 +238,23 @@ 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) { + // 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 305ef99f..0f65a214 100644 --- a/e2e/tests/chat.spec.ts +++ b/e2e/tests/chat.spec.ts @@ -87,11 +87,18 @@ test.describe('Chat Feature Tests', () => { const processingComplete = await homePage.waitForProcessingToComplete(); expect(processingComplete).toBeTruthy(); - // Verify Query Analysis message appears (but without actual SQL) - const sqlMessageVisible = await homePage.isSQLQueryMessageVisible(); - expect(sqlMessageVisible).toBeTruthy(); - - // Verify NO actual SQL content (should say "Query Analysis" or "Off topic") + // 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. + // 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"); expect(hasSQLContent).toBeFalsy(); 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_db_execution_timeouts.py b/tests/test_db_execution_timeouts.py new file mode 100644 index 00000000..45786bbe --- /dev/null +++ b/tests/test_db_execution_timeouts.py @@ -0,0 +1,341 @@ +"""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 + +from api.config import Config +# 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" + + +@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_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 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 +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): + 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 + ) + + +@pytest.mark.unit +@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, expected_ms, reason): + """Exactly one directive survives, never looser than the 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}") + 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("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 + ("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 +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}" + ) + + + +@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 + + +@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..24841ef9 --- /dev/null +++ b/tests/test_deadline_guard.py @@ -0,0 +1,78 @@ +"""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 + +# One import style: the module object is needed anyway, because the tests +# monkeypatch its grace constant. +from api.loaders import deadline + + +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.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.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.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.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 new file mode 100644 index 00000000..47318445 --- /dev/null +++ b/tests/test_embeddings_offloading.py @@ -0,0 +1,144 @@ +"""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.""" + stop = asyncio.Event() + + async def ticker(): + samples = [] + while not stop.is_set(): + samples.append(time.monotonic()) + await asyncio.sleep(TICK) + return samples + + ticker_task = asyncio.ensure_future(ticker()) + result = await coro + stop.set() + return result, await ticker_task + + +@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_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.""" + kwargs = Config.EMBEDDING_MODEL._embedding_kwargs() + # 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 + assert kwargs["max_retries"] == 0 + assert kwargs["num_retries"] == 0 diff --git a/tests/test_find_offloading.py b/tests/test_find_offloading.py new file mode 100644 index 00000000..5ba1b888 --- /dev/null +++ b/tests/test_find_offloading.py @@ -0,0 +1,102 @@ +"""``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.""" + stop = asyncio.Event() + + async def ticker(): + samples = [] + while not stop.is_set(): + 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() + 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)" + + # 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" + + # The fake graph returns no rows, so the call still completes normally. + assert tables == [] diff --git a/tests/test_loader_contract.py b/tests/test_loader_contract.py new file mode 100644 index 00000000..ecdce4f0 --- /dev/null +++ b/tests/test_loader_contract.py @@ -0,0 +1,51 @@ +"""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 + +# 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] + + +@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 new file mode 100644 index 00000000..0e056cb3 --- /dev/null +++ b/tests/test_schema_load_offloading.py @@ -0,0 +1,383 @@ +"""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 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 +TICK = 0.02 + + +async def _ticks_while_consuming(agen): + """Drain *agen*, counting event-loop ticks.""" + stop = asyncio.Event() + + async def ticker(): + samples = [] + while not stop.is_set(): + 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() + return steps, await ticker_task + + +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" + + +@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: + # Expected: we cancelled it. What matters is the cleanup that runs + # afterwards, asserted below. + 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" + + + +@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, "_EXECUTOR", 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" + + +@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, "_EXECUTOR", 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) + + + +@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 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 diff --git a/tests/test_stream_idle_timeout.py b/tests/test_stream_idle_timeout.py new file mode 100644 index 00000000..3ae58b62 --- /dev/null +++ b/tests/test_stream_idle_timeout.py @@ -0,0 +1,249 @@ +"""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.""" + from api.core import 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) + + +@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) diff --git a/tests/test_stream_keepalive.py b/tests/test_stream_keepalive.py new file mode 100644 index 00000000..e0466150 --- /dev/null +++ b/tests/test_stream_keepalive.py @@ -0,0 +1,200 @@ +"""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) + + +@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: + # 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}") + + 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) + + +@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" diff --git a/tests/test_timeout_validation.py b/tests/test_timeout_validation.py new file mode 100644 index 00000000..95665d56 --- /dev/null +++ b/tests/test_timeout_validation.py @@ -0,0 +1,181 @@ +"""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 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) + # ``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 +@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 + + +@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, across every attempt. + + 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 import config as api_config + + monkeypatch.setattr(api_config.Config, "LLM_TIMEOUT", 90.0, raising=False) + monkeypatch.setattr(api_config.Config, "LLM_MAX_RETRIES", retries, raising=False) + + 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 + # Default is the whole budget, which is right for single-attempt callers. + assert bounds["timeout"] == 90.0 + + +@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 + + +@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}