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
|
"""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).
|
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
|
TODO(follow-up, PR #814 review): thread etags for file / deck_card /
|
||||||
(only note scans pass etag today). Until then their JetStream Nats-Msg-Id
|
news_item scans too (only note scans pass etag today). Until then their
|
||||||
dedup keys off modified_at, which misses content changes that leave
|
JetStream Nats-Msg-Id dedup keys off modified_at, which misses content
|
||||||
modified_at unchanged (e.g. a file move/rename).
|
changes that leave modified_at unchanged (e.g. a file move/rename).
|
||||||
"""
|
"""
|
||||||
return task.etag or str(task.modified_at)
|
return task.etag or str(task.modified_at)
|
||||||
|
|
||||||
|
|||||||
@@ -144,12 +144,29 @@ class NatsStatusSubscriber:
|
|||||||
import nats.errors # noqa: PLC0415
|
import nats.errors # noqa: PLC0415
|
||||||
|
|
||||||
subject = f"mcp.document.*.{self.tenant_id}"
|
subject = f"mcp.document.*.{self.tenant_id}"
|
||||||
sub = await self._js.pull_subscribe(
|
# Signal "task running" *before* the first (fallible) subscribe: bus
|
||||||
subject, durable=f"mcp-status-{self.tenant_id}"
|
# 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:
|
if task_status is not None:
|
||||||
task_status.started()
|
task_status.started()
|
||||||
|
|
||||||
|
sub = None
|
||||||
while not shutdown_event.is_set():
|
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:
|
try:
|
||||||
msgs = await sub.fetch(batch=16, timeout=5)
|
msgs = await sub.fetch(batch=16, timeout=5)
|
||||||
except nats.errors.TimeoutError:
|
except nats.errors.TimeoutError:
|
||||||
@@ -157,11 +174,14 @@ class NatsStatusSubscriber:
|
|||||||
# straight back to re-check shutdown — no log, no extra sleep.
|
# straight back to re-check shutdown — no log, no extra sleep.
|
||||||
continue
|
continue
|
||||||
except Exception:
|
except Exception:
|
||||||
# Real broker error (disconnect, auth failure, stream deleted).
|
# Real broker error (disconnect, auth failure, stream deleted):
|
||||||
# Log it and back off so we don't hot-spin against a dead broker.
|
# drop the (possibly dead) subscription, back off, and
|
||||||
|
# re-subscribe on the next iteration rather than hot-spinning.
|
||||||
logger.warning(
|
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)
|
await anyio.sleep(5)
|
||||||
continue
|
continue
|
||||||
for msg in msgs:
|
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.document.ready.t1", b"not json")
|
||||||
sub.handle_message("mcp.ingest.requested.t1", b'{"doc_id":"x"}')
|
sub.handle_message("mcp.ingest.requested.t1", b'{"doc_id":"x"}')
|
||||||
assert len(store) == 0
|
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