fix(db_utils): make acquire_with_retry yield exactly once (#1880)

acquire_with_retry's retry loop wrapped the yield, violating
@asynccontextmanager's single-yield contract. When user code inside
the async with block raised a retryable exception, the loop iterated
and tried to yield again, producing RuntimeError("generator didn't
stop after athrow()") on every retryable inner error. This masked
the real cause and was the root of 1,934 identical failed
consolidation ops on shurick-memory in production since 2026-03-30.

Retry now wraps only the acquire (via AsyncExitStack). User-code
exceptions inside the block propagate as their real types — strictly
better for observability, since the prior retry-of-user-code branch
was already non-functional (always crashed with the RuntimeError above).

Includes a regression unit test asserting (a) the original retryable
exception propagates unchanged and (b) the connection is released
exactly once.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: slayoffer

hindsight-api-slim/hindsight_api/engine/db_utils.py

@@ -6,7 +6,7 @@
 import logging
 import time
 from collections.abc import AsyncIterator
-from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack, asynccontextmanager
 from typing import Any
 
 logger = logging.getLogger(__name__)
@@ -101,6 +101,14 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
     """
     Async context manager to acquire a database connection with retry logic.
 
+    Retries the *acquire* itself when it raises a retryable error (connection
+    drop, timeout, deadlock detected during acquire). Exceptions raised by
+    user code inside the ``async with`` block are NOT retried — they propagate
+    as-is. Wrapping retry around the yield would violate the
+    ``@asynccontextmanager`` single-yield contract and surface as
+    ``RuntimeError("generator didn't stop after athrow()")`` on every
+    retryable inner error, masking the real cause.
+
     Accepts either a DatabaseBackend or a raw asyncpg.Pool for backward compatibility.
 
     Usage:
@@ -109,39 +117,40 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
 
     Args:
         backend_or_pool: A DatabaseBackend instance or asyncpg.Pool
-        max_retries: Maximum number of retry attempts
+        max_retries: Maximum number of retry attempts for the acquire step
 
     Yields:
         A DatabaseConnection (if backend) or asyncpg.Connection (if pool)
     """
     from .db.base import DatabaseBackend
 
     if isinstance(backend_or_pool, DatabaseBackend) or getattr(backend_or_pool, "_wraps_backend", False):
-        # Use the backend's acquire context manager with retry
         start = time.time()
-        last_exception = None
-        for attempt in range(max_retries + 1):
-            try:
-                async with backend_or_pool.acquire() as conn:
-                    acquire_time = time.time() - start
-                    if acquire_time > 0.05:
-                        logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
-                    yield conn
-                    return
-            except Exception as e:
-                if not _is_retryable(e):
-                    raise
-                last_exception = e
-                if attempt < max_retries:
-                    delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
-                    logger.warning(
-                        f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
-                        f"Retrying in {delay:.1f}s..."
-                    )
-                    await asyncio.sleep(delay)
-                else:
-                    logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
-        raise last_exception
+        async with AsyncExitStack() as stack:
+            conn: Any = None
+            for attempt in range(max_retries + 1):
+                try:
+                    conn = await stack.enter_async_context(backend_or_pool.acquire())
+                    break
+                except Exception as e:
+                    if not _is_retryable(e):
+                        raise
+                    if attempt < max_retries:
+                        delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
+                        logger.warning(
+                            f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
+                            f"Retrying in {delay:.1f}s..."
+                        )
+                        await asyncio.sleep(delay)
+                    else:
+                        logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
+                        raise
+
+            acquire_time = time.time() - start
+            if acquire_time > 0.05:
+                logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
+
+            yield conn
     else:
         # Legacy path: raw asyncpg.Pool
         pool = backend_or_pool

hindsight-api-slim/tests/test_db_utils.py

@@ -0,0 +1,79 @@
+"""Tests for hindsight_api.engine.db_utils.
+
+Regression coverage for the single-yield contract of ``acquire_with_retry`` —
+historically a retry loop wrapped the ``yield`` and caused every retryable
+user-code exception to surface as
+``RuntimeError("generator didn't stop after athrow()")``, masking the real
+cause and producing identical failed-op rows in production (see the 1,934
+failed consolidations on ``shurick-memory`` in May 2026).
+"""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager
+
+import pytest
+
+from hindsight_api.engine.db_utils import acquire_with_retry
+
+
+class _FakeConnection:
+    """Stand-in for a DatabaseConnection that records release events."""
+
+    def __init__(self) -> None:
+        self.released = 0
+
+
+class _FakeBackend:
+    """Duck-typed DatabaseBackend that opts in via the ``_wraps_backend`` flag.
+
+    ``acquire_with_retry`` accepts either a real ``DatabaseBackend`` subclass
+    or any object with ``_wraps_backend = True``; the flag avoids having to
+    stub the full abstract surface for unit tests.
+    """
+
+    _wraps_backend = True
+
+    def __init__(self) -> None:
+        self.acquired = 0
+        self.last_conn: _FakeConnection | None = None
+
+    @asynccontextmanager
+    async def acquire(self):
+        self.acquired += 1
+        conn = _FakeConnection()
+        self.last_conn = conn
+        try:
+            yield conn
+        finally:
+            conn.released += 1
+
+
+@pytest.mark.asyncio
+async def test_retryable_user_code_exception_propagates_unchanged():
+    """A retryable exception inside ``async with`` must propagate as itself.
+
+    Before the single-yield refactor, the retry loop around the ``yield``
+    re-entered ``yield conn`` on the next iteration, violating
+    ``@asynccontextmanager``'s contract and surfacing as
+    ``RuntimeError("generator didn't stop after athrow()")`` — the symptom
+    that broke consolidation on large banks.
+    """
+
+    backend = _FakeBackend()
+    sentinel = asyncio.TimeoutError("query exceeded statement_timeout")
+
+    with pytest.raises(asyncio.TimeoutError) as excinfo:
+        async with acquire_with_retry(backend) as conn:
+            assert isinstance(conn, _FakeConnection)
+            raise sentinel
+
+    # The original exception flows out — not a RuntimeError wrapper.
+    assert excinfo.value is sentinel
+    assert not isinstance(excinfo.value, RuntimeError)
+
+    # Acquire was called exactly once — user-code failure must not retry.
+    assert backend.acquired == 1
+    assert backend.last_conn is not None
+    assert backend.last_conn.released == 1, "connection must be released exactly once"