mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): propagate CancelledError from memory init cleanly
External cancellation while _ensure_initialized is in flight raises CancelledError (BaseException). The existing 'except asyncio.TimeoutError:' branch does not catch it, and caller 'except Exception:' blocks don't either, so it propagated as-is with _backend possibly still assigned. Add an explicit 'except asyncio.CancelledError:' handler that nulls _backend, clears _initialized, logs at info, and re-raises. Cancellation is a shutdown signal, not an error to swallow — but the state must be clean so a later retry starts fresh.
This commit is contained in:
parent
3f1f82d7e2
commit
196d8d69a7
2 changed files with 53 additions and 0 deletions
|
|
@ -198,6 +198,20 @@ class MemoryHandler:
|
|||
"Subsequent requests will retry."
|
||||
)
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
# External cancellation (shutdown / task cancelled).
|
||||
# CancelledError is BaseException — the TimeoutError branch
|
||||
# above does NOT catch it, and caller ``except Exception``
|
||||
# blocks don't either, so it propagates unconditionally.
|
||||
# Reset state so any later retry starts clean, then re-raise:
|
||||
# cancellation is a signal, not an error to swallow.
|
||||
self._backend = None
|
||||
self._initialized = False
|
||||
logger.info(
|
||||
"Memory: backend initialization cancelled "
|
||||
f"(backend={self.config.backend})"
|
||||
)
|
||||
raise
|
||||
|
||||
async def _init_backend_locked(self) -> None:
|
||||
"""Actual backend-init body. Must be called with ``_init_lock`` held."""
|
||||
|
|
|
|||
|
|
@ -190,6 +190,45 @@ async def test_ensure_initialized_timeout_nulls_partially_initialized_backend(
|
|||
assert handler._backend is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_initialized_cancellation_propagates_and_resets_state(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
"""External cancellation of an in-flight ``_ensure_initialized`` must
|
||||
propagate (CancelledError is BaseException — not a swallowable error)
|
||||
and leave the handler in a clean state."""
|
||||
|
||||
class HangingBackend:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
async def _ensure_initialized(self) -> None:
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
|
||||
import headroom.memory.backends.local as local_mod
|
||||
|
||||
monkeypatch.setattr(local_mod, "LocalBackend", HangingBackend)
|
||||
|
||||
handler = MemoryHandler(
|
||||
MemoryConfig(
|
||||
enabled=True, backend="local", db_path=str(tmp_path / "mem.db")
|
||||
)
|
||||
)
|
||||
|
||||
task = asyncio.create_task(handler._ensure_initialized())
|
||||
# Give the task a tick to enter _init_backend_locked and assign _backend.
|
||||
await asyncio.sleep(0.05)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
assert handler._initialized is False
|
||||
assert handler._backend is None
|
||||
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Real backend init (no monkeypatching) — integration smoke test
|
||||
# -------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue