fix(vector): make NATS status subscriber resilient at startup
Follow-up to PR #814 review. NatsStatusSubscriber.run() called task_status.started() *after* the fallible pull_subscribe, so a NATS broker that wasn't ready when the MCP server started would crash the lifespan instead of retrying. Bus status is a non-critical observability path, so: - signal started() before the first subscribe (semantics: "loop is running", not "subscription succeeded"); - retry a failed subscribe with backoff instead of propagating; - on a real fetch error (not an idle timeout) drop the subscription and re-subscribe rather than fetching against a possibly-dead handle. Also anchor the _content_hash etag-threading TODO to the PR #814 review thread so it is discoverable outside git blame. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
f65ebfc65b
commit
cd2145df09
@@ -54,10 +54,10 @@ def _content_hash(task: DocumentTask) -> str:
|
||||
"""etag is the change-detection token; fall back to modified_at when it is
|
||||
absent (e.g. deletes, or sources whose etag we don't thread through).
|
||||
|
||||
TODO(follow-up): thread etags for file / deck_card / news_item scans too
|
||||
(only note scans pass etag today). Until then their JetStream Nats-Msg-Id
|
||||
dedup keys off modified_at, which misses content changes that leave
|
||||
modified_at unchanged (e.g. a file move/rename).
|
||||
TODO(follow-up, PR #814 review): thread etags for file / deck_card /
|
||||
news_item scans too (only note scans pass etag today). Until then their
|
||||
JetStream Nats-Msg-Id dedup keys off modified_at, which misses content
|
||||
changes that leave modified_at unchanged (e.g. a file move/rename).
|
||||
"""
|
||||
return task.etag or str(task.modified_at)
|
||||
|
||||
|
||||
@@ -144,12 +144,29 @@ class NatsStatusSubscriber:
|
||||
import nats.errors # noqa: PLC0415
|
||||
|
||||
subject = f"mcp.document.*.{self.tenant_id}"
|
||||
sub = await self._js.pull_subscribe(
|
||||
subject, durable=f"mcp-status-{self.tenant_id}"
|
||||
)
|
||||
# Signal "task running" *before* the first (fallible) subscribe: bus
|
||||
# status is a non-critical observability path, so a broker that isn't
|
||||
# ready at startup should retry below rather than crash the lifespan.
|
||||
# ``started()`` therefore means "the subscriber loop is running", not
|
||||
# "the subscription succeeded".
|
||||
if task_status is not None:
|
||||
task_status.started()
|
||||
|
||||
sub = None
|
||||
while not shutdown_event.is_set():
|
||||
if sub is None:
|
||||
try:
|
||||
sub = await self._js.pull_subscribe(
|
||||
subject, durable=f"mcp-status-{self.tenant_id}"
|
||||
)
|
||||
except Exception:
|
||||
# Broker not ready / transient connect error: back off and
|
||||
# retry the subscribe instead of giving up.
|
||||
logger.warning(
|
||||
"NATS status subscribe failed; retrying", exc_info=True
|
||||
)
|
||||
await anyio.sleep(5)
|
||||
continue
|
||||
try:
|
||||
msgs = await sub.fetch(batch=16, timeout=5)
|
||||
except nats.errors.TimeoutError:
|
||||
@@ -157,11 +174,14 @@ class NatsStatusSubscriber:
|
||||
# straight back to re-check shutdown — no log, no extra sleep.
|
||||
continue
|
||||
except Exception:
|
||||
# Real broker error (disconnect, auth failure, stream deleted).
|
||||
# Log it and back off so we don't hot-spin against a dead broker.
|
||||
# Real broker error (disconnect, auth failure, stream deleted):
|
||||
# drop the (possibly dead) subscription, back off, and
|
||||
# re-subscribe on the next iteration rather than hot-spinning.
|
||||
logger.warning(
|
||||
"NATS status subscriber fetch failed; retrying", exc_info=True
|
||||
"NATS status subscriber fetch failed; re-subscribing",
|
||||
exc_info=True,
|
||||
)
|
||||
sub = None
|
||||
await anyio.sleep(5)
|
||||
continue
|
||||
for msg in msgs:
|
||||
|
||||
@@ -68,3 +68,100 @@ def test_handle_message_ignores_bad_payload_and_subject():
|
||||
sub.handle_message("mcp.document.ready.t1", b"not json")
|
||||
sub.handle_message("mcp.ingest.requested.t1", b'{"doc_id":"x"}')
|
||||
assert len(store) == 0
|
||||
|
||||
|
||||
async def test_run_signals_started_then_retries_subscribe(mocker, monkeypatch):
|
||||
"""run() signals started before subscribing, retries a failed subscribe,
|
||||
and consumes messages once subscribed."""
|
||||
import anyio
|
||||
|
||||
# Make backoff sleeps instant so the retry path doesn't stall the test.
|
||||
async def _no_sleep(*_a, **_k):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(anyio, "sleep", _no_sleep)
|
||||
|
||||
store = StatusStore()
|
||||
js = mocker.AsyncMock()
|
||||
|
||||
# First subscribe attempt fails (broker not ready), second succeeds.
|
||||
fake_sub = mocker.AsyncMock()
|
||||
js.pull_subscribe.side_effect = [ConnectionError("broker not ready"), fake_sub]
|
||||
|
||||
shutdown = anyio.Event()
|
||||
msg = mocker.Mock()
|
||||
msg.subject = "mcp.document.ready.t1"
|
||||
msg.data = json.dumps({"doc_id": "d1", "content_hash": "h1"}).encode()
|
||||
msg.ack = mocker.AsyncMock()
|
||||
|
||||
fetches = {"n": 0}
|
||||
|
||||
async def _fetch(*_a, **_k):
|
||||
fetches["n"] += 1
|
||||
if fetches["n"] == 1:
|
||||
return [msg]
|
||||
shutdown.set() # stop the loop after the first batch is handled
|
||||
return []
|
||||
|
||||
fake_sub.fetch.side_effect = _fetch
|
||||
|
||||
task_status = mocker.Mock()
|
||||
subscriber = NatsStatusSubscriber(
|
||||
nc=mocker.AsyncMock(), js=js, tenant_id="t1", store=store
|
||||
)
|
||||
|
||||
await subscriber.run(shutdown, task_status=task_status)
|
||||
|
||||
# started() fires before any subscribe attempt and exactly once.
|
||||
task_status.started.assert_called_once()
|
||||
# The failed first subscribe was retried (two attempts total).
|
||||
assert js.pull_subscribe.call_count == 2
|
||||
# The message from the successful subscription was recorded + acked.
|
||||
assert store.get("d1") == {
|
||||
"state": "ready",
|
||||
"content_hash": "h1",
|
||||
"transitioned_at": None,
|
||||
}
|
||||
msg.ack.assert_awaited_once()
|
||||
|
||||
|
||||
async def test_run_resubscribes_after_fetch_error(mocker, monkeypatch):
|
||||
"""A non-timeout fetch error drops the subscription and re-subscribes."""
|
||||
import anyio
|
||||
import nats.errors
|
||||
|
||||
async def _no_sleep(*_a, **_k):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(anyio, "sleep", _no_sleep)
|
||||
|
||||
store = StatusStore()
|
||||
js = mocker.AsyncMock()
|
||||
first_sub = mocker.AsyncMock()
|
||||
second_sub = mocker.AsyncMock()
|
||||
js.pull_subscribe.side_effect = [first_sub, second_sub]
|
||||
|
||||
shutdown = anyio.Event()
|
||||
|
||||
# first_sub.fetch raises a real broker error → re-subscribe.
|
||||
first_sub.fetch.side_effect = ConnectionResetError("broker dropped")
|
||||
|
||||
# second_sub.fetch idles once (timeout) then stops the loop.
|
||||
fetches = {"n": 0}
|
||||
|
||||
async def _second_fetch(*_a, **_k):
|
||||
fetches["n"] += 1
|
||||
if fetches["n"] == 1:
|
||||
raise nats.errors.TimeoutError
|
||||
shutdown.set()
|
||||
return []
|
||||
|
||||
second_sub.fetch.side_effect = _second_fetch
|
||||
|
||||
subscriber = NatsStatusSubscriber(
|
||||
nc=mocker.AsyncMock(), js=js, tenant_id="t1", store=store
|
||||
)
|
||||
await subscriber.run(shutdown)
|
||||
|
||||
# Re-subscribed after the fetch error (two subscriptions used).
|
||||
assert js.pull_subscribe.call_count == 2
|
||||
|
||||
Reference in New Issue
Block a user