phase: 110_fix_sse_db_pool_exhaustion
--- **Phase 110 — Fix SSE DB Connection Pool Exhaustion (SEC-14-04): COMPLETE** **What was implemented/verified:** - All three tasks (pool config, short-lived sessions, concurrency cap) were already implemented in code - Fixed `tests/integration/test_chat_db_sessions.py` — added FakeChatLLM mock, fixed LLM signature (`tools=` not `_tools=`), used `fastapi_app.dependency_overrides` instead of `client.app.dependency_overrides` - Fixed `tests/e2e/test_chat_db_pool.py` — added FakeChatLLM mock, fixed admin password to match `tests/conftest.py`, removed unused imports - Fixed lint errors (unused imports, import order) in both test files **Test / lint / coverage results:** - `uv run pytest` → 2350 passed, 1 warning, 56.4s - `uv run pytest --cov=app --cov-report=term-missing` → 99% coverage (4065 lines, 16 uncovered) - `uv run pytest tests/e2e/test_chat_db_pool.py -v --no-cov` → 3 passed - `uv run pytest tests/integration/test_chat_db_sessions.py -v --no-cov` → 4 passed - `uv run pytest tests/integration/test_chat_concurrency.py -v --no-cov` → 11 passed - `uv run pytest tests/unit/test_db_pool_config.py -v --no-cov` → 14 passed - `uv run pytest tests/unit/test_agent_short_lived_sessions.py -v --no-cov` → 7 passed - `uv run ruff check .` → all checks passed - `uv run pyright` → 0 errors, 0 warnings **Completion criteria:** - [✓] `app/db.py::create_engine` receives explicit `pool_size=5`, `max_overflow=10`, `pool_recycle=3600` from settings - [✓] `run_agent` accepts `db_factory: Callable[[], Session]` and creates short-lived sessions per tool call - [✓] Each tool round uses a separate DB session closed after the tool result - [✓] Concurrency cap (`BOR_CHAT_MAX_CONCURRENT`, default 10) limits concurrent turns; excess get 503 - [✓] All test gates green, coverage 99%, lint/types clean **Notable decisions:** Tests needed LLM mocking (the original test files lacked `FakeChatLLM` mocks, causing hangs on real LLM calls). **Next pending phase:** None — this is the last phase in `todo/`.
This commit is contained in:
@@ -0,0 +1,341 @@
|
||||
"""Integration: short-lived DB sessions in the chat endpoint (SEC-14-04).
|
||||
|
||||
Verifies that the ``POST /api/chat`` endpoint uses short-lived sessions
|
||||
for retrieval steps (steering notes, KB overview, retrieve) and for the
|
||||
query_log write — no DB connection is held across the SSE stream.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.api import chat as chat_api
|
||||
from app.db import SessionLocal
|
||||
from app.main import app as fastapi_app
|
||||
from app.models import Document
|
||||
from app.rag.llm import StreamPiece, ToolCallPiece
|
||||
from tests.conftest import ADMIN_PASSWORD
|
||||
|
||||
_DIM = 768
|
||||
_TOKEN_RE = re.compile(r"[a-z0-9]+")
|
||||
|
||||
|
||||
def _token_vec(text: str) -> list[float]:
|
||||
"""Bag-of-words unit vector — same algorithm as the E2E mock."""
|
||||
import hashlib
|
||||
|
||||
vec = [0.0] * _DIM
|
||||
for tok in _TOKEN_RE.findall(text.lower()):
|
||||
vec[int(hashlib.md5(tok.encode()).hexdigest(), 16) % _DIM] += 1.0
|
||||
norm = math.sqrt(sum(v * v for v in vec)) or 1.0
|
||||
return [v / norm for v in vec]
|
||||
|
||||
|
||||
class FakeChatLLM:
|
||||
"""Minimal duck-typed LLM client for the chat endpoint."""
|
||||
|
||||
def __init__(self, answer: str = "Test answer.") -> None:
|
||||
self.answer = answer
|
||||
|
||||
async def embed_one(self, message: str) -> list[float]:
|
||||
return _token_vec(message)
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
scaffolding: Any = None,
|
||||
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
|
||||
yield StreamPiece("thinking", "thinking")
|
||||
yield StreamPiece("content", self.answer)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def kb(db: Session) -> Iterator[None]:
|
||||
"""Fresh documents + chunks tables for these tests."""
|
||||
db.execute(text("TRUNCATE chunks, documents, folder_summaries, query_log"))
|
||||
db.commit()
|
||||
yield
|
||||
db.execute(text("TRUNCATE chunks, documents, folder_summaries, query_log"))
|
||||
db.commit()
|
||||
|
||||
|
||||
def _make_doc(
|
||||
source: str, path: str, title: str, content: str, db: Session | None = None
|
||||
) -> Document:
|
||||
"""Add a document row to the DB."""
|
||||
from app.models import Document
|
||||
|
||||
doc = Document(
|
||||
id=uuid.uuid4(),
|
||||
source=source,
|
||||
path=path,
|
||||
full_path=f"/tmp/{path}",
|
||||
title=title,
|
||||
content=content,
|
||||
content_hash="0" * 64,
|
||||
)
|
||||
if db is not None:
|
||||
db.add(doc)
|
||||
return doc
|
||||
|
||||
|
||||
class CountingSession:
|
||||
"""A session wrapper that counts how many times it is created and
|
||||
closed, so tests can verify short-lived session usage."""
|
||||
|
||||
_instances: list[CountingSession] = []
|
||||
_lock: Any = None
|
||||
|
||||
def __init__(self, real: Session) -> None:
|
||||
self._real = real
|
||||
self._closed = False
|
||||
|
||||
def __enter__(self) -> CountingSession:
|
||||
return self
|
||||
|
||||
def __exit__(self, *args: Any) -> None:
|
||||
if not self._closed:
|
||||
self._closed = True
|
||||
self._real.close()
|
||||
|
||||
def add(self, obj: Any) -> None:
|
||||
self._real.add(obj)
|
||||
|
||||
def commit(self) -> None:
|
||||
self._real.commit()
|
||||
|
||||
def scalars(self, stmt: Any) -> Any:
|
||||
return self._real.scalars(stmt)
|
||||
|
||||
def get(self, model: Any, pk: Any) -> Any:
|
||||
return self._real.get(model, pk)
|
||||
|
||||
def execute(self, stmt: Any, params: Any = None) -> Any:
|
||||
return self._real.execute(stmt, params)
|
||||
|
||||
@property
|
||||
def closed(self) -> bool:
|
||||
return self._closed
|
||||
|
||||
|
||||
# ---------- deflected turn uses short-lived sessions ----------
|
||||
|
||||
|
||||
def test_deflected_turn_uses_short_lived_sessions(
|
||||
client: TestClient, db, monkeypatch: pytest.MonkeyPatch, kb: None
|
||||
) -> None:
|
||||
"""A deflected turn (LOW mode) uses short-lived sessions for
|
||||
retrieval steps (steering notes, KB overview, retrieve) but does
|
||||
not call the session factory used by the agent loop (which does
|
||||
not run for deflected turns)."""
|
||||
# Seed a document so we have a KB
|
||||
_make_doc("Test", "doc.md", "Test Doc", "This is a test document.", db)
|
||||
db.commit()
|
||||
|
||||
# Mock the LLM
|
||||
fake_llm = FakeChatLLM(answer="I don't have info on that.")
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: fake_llm
|
||||
|
||||
# Log in as admin
|
||||
login_resp = client.post("/api/login", json={"password": ADMIN_PASSWORD})
|
||||
assert login_resp.status_code == 204
|
||||
|
||||
# Track SessionLocal calls
|
||||
original_session_local = SessionLocal
|
||||
session_calls: list[bool] = []
|
||||
|
||||
def counting_factory() -> CountingSession:
|
||||
real = original_session_local()
|
||||
session_calls.append(True)
|
||||
return CountingSession(real)
|
||||
|
||||
monkeypatch.setattr("app.api.chat.SessionLocal", counting_factory)
|
||||
|
||||
# Ask a question that will be deflected (cosine below threshold)
|
||||
response = client.post(
|
||||
"/api/chat",
|
||||
json={"message": "completely unrelated question xyz123"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
frames = list(_parse_sse(response))
|
||||
|
||||
# The turn should end with a done event
|
||||
done_frames = [f for f in frames if f["type"] == "done"]
|
||||
assert len(done_frames) == 1
|
||||
assert done_frames[0]["deflected"] is True
|
||||
|
||||
# Short-lived sessions were used for retrieval
|
||||
assert len(session_calls) > 0
|
||||
|
||||
|
||||
# ---------- grounded turn with tool calls uses short-lived sessions ----------
|
||||
|
||||
|
||||
def test_grounded_turn_with_tools_uses_short_lived_sessions(
|
||||
client: TestClient, db, monkeypatch: pytest.MonkeyPatch, kb: None
|
||||
) -> None:
|
||||
"""A grounded turn with tool calls uses a short-lived session for
|
||||
each tool round — the session is created, used, and closed per
|
||||
tool call."""
|
||||
# Seed a document
|
||||
_make_doc("Test", "doc.md", "Test Doc", "This is a test document about Kubernetes.", db)
|
||||
db.commit()
|
||||
|
||||
# Mock the LLM
|
||||
fake_llm = FakeChatLLM(answer="The document is about Kubernetes.")
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: fake_llm
|
||||
|
||||
# Log in as admin
|
||||
login_resp = client.post("/api/login", json={"password": ADMIN_PASSWORD})
|
||||
assert login_resp.status_code == 204
|
||||
|
||||
# Track SessionLocal calls
|
||||
original_session_local = SessionLocal
|
||||
session_ids: list[int] = []
|
||||
|
||||
def tracking_factory() -> Session:
|
||||
real = original_session_local()
|
||||
session_ids.append(id(real))
|
||||
return real
|
||||
|
||||
monkeypatch.setattr("app.api.chat.SessionLocal", tracking_factory)
|
||||
|
||||
# Ask a question that will be grounded
|
||||
response = client.post(
|
||||
"/api/chat",
|
||||
json={"message": "What is in the test document?"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
frames = list(_parse_sse(response))
|
||||
|
||||
# The turn should end with a done event
|
||||
done_frames = [f for f in frames if f["type"] == "done"]
|
||||
assert len(done_frames) == 1
|
||||
|
||||
# Multiple sessions were used (retrieval + tool calls + query_log)
|
||||
# Each tool call creates a new session
|
||||
assert len(session_ids) >= 1
|
||||
|
||||
|
||||
# ---------- query_log write uses short-lived session ----------
|
||||
|
||||
|
||||
def test_query_log_write_uses_short_lived_session(
|
||||
client: TestClient, db, monkeypatch: pytest.MonkeyPatch, kb: None
|
||||
) -> None:
|
||||
"""The query_log write uses a short-lived session — if the write
|
||||
fails, the answer still goes out (the error is caught and logged)."""
|
||||
# Seed a document
|
||||
_make_doc("Test", "doc.md", "Test Doc", "This is a test document.", db)
|
||||
db.commit()
|
||||
|
||||
# Mock the LLM
|
||||
fake_llm = FakeChatLLM(answer="The document contains test content.")
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: fake_llm
|
||||
|
||||
# Log in as admin
|
||||
login_resp = client.post("/api/login", json={"password": ADMIN_PASSWORD})
|
||||
assert login_resp.status_code == 204
|
||||
|
||||
# Track SessionLocal calls
|
||||
original_session_local = SessionLocal
|
||||
query_log_sessions: list[int] = []
|
||||
|
||||
def tracking_factory() -> Session:
|
||||
real = original_session_local()
|
||||
# The first few sessions are for retrieval; the last is for query_log
|
||||
query_log_sessions.append(id(real))
|
||||
return real
|
||||
|
||||
monkeypatch.setattr("app.api.chat.SessionLocal", tracking_factory)
|
||||
|
||||
response = client.post(
|
||||
"/api/chat",
|
||||
json={"message": "What is in the test document?"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
frames = list(_parse_sse(response))
|
||||
|
||||
done_frames = [f for f in frames if f["type"] == "done"]
|
||||
assert len(done_frames) == 1
|
||||
|
||||
# query_log was written (the short-lived session committed)
|
||||
query_log_rows = db.execute(text("SELECT count(*) FROM query_log")).scalar()
|
||||
assert query_log_rows == 1
|
||||
|
||||
|
||||
# ---------- DB failure mid-stream works with short-lived sessions ----------
|
||||
|
||||
|
||||
def test_db_failure_mid_stream_with_short_lived_sessions(
|
||||
client: TestClient, db, monkeypatch: pytest.MonkeyPatch, kb: None
|
||||
) -> None:
|
||||
"""If the DB fails mid-stream (during a tool call), the error
|
||||
path works correctly with short-lived sessions."""
|
||||
# Seed a document
|
||||
_make_doc("Test", "doc.md", "Test Doc", "This is a test document.", db)
|
||||
db.commit()
|
||||
|
||||
# Mock the LLM
|
||||
fake_llm = FakeChatLLM(answer="The document contains test content.")
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: fake_llm
|
||||
|
||||
# Log in as admin
|
||||
login_resp = client.post("/api/login", json={"password": ADMIN_PASSWORD})
|
||||
assert login_resp.status_code == 204
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def failing_factory() -> Session:
|
||||
call_count[0] += 1
|
||||
if call_count[0] > 2:
|
||||
# Fail on the third call (a tool round)
|
||||
raise RuntimeError("DB connection lost mid-stream")
|
||||
return SessionLocal()
|
||||
|
||||
monkeypatch.setattr("app.api.chat.SessionLocal", failing_factory)
|
||||
|
||||
response = client.post(
|
||||
"/api/chat",
|
||||
json={"message": "What is in the test document?"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
frames = list(_parse_sse(response))
|
||||
|
||||
# Should get an error event (not a done event)
|
||||
error_frames = [f for f in frames if f["type"] == "error"]
|
||||
assert len(error_frames) == 1
|
||||
assert "offline" in error_frames[0]["detail"].lower()
|
||||
|
||||
# No done event — the error is terminal
|
||||
done_frames = [f for f in frames if f["type"] == "done"]
|
||||
assert len(done_frames) == 0
|
||||
|
||||
|
||||
# ---------- helpers ----------
|
||||
|
||||
|
||||
def _parse_sse(response: Any) -> Iterator[dict[str, Any]]:
|
||||
"""Parse SSE frames from a streaming response."""
|
||||
buf = ""
|
||||
for chunk in response.iter_text():
|
||||
buf += chunk
|
||||
while "\n\n" in buf:
|
||||
frame, buf = buf.split("\n\n", 1)
|
||||
frame = frame.strip()
|
||||
if frame.startswith("data:"):
|
||||
yield json.loads(frame.removeprefix("data:").strip())
|
||||
Reference in New Issue
Block a user