feat: add opt-in MCP decomposition hook points (design §10)

Adds the seven §10.2 hook-point modules + five env vars so Astrolabe Cloud can
offload document processing to the external document-processor / embedding
gateway. Purely additive: with every setting unset the server behaves exactly
as today, so self-hosters are unaffected (Deck #92).

Hook points (all default to current monolith behavior):
- config: EMBEDDING_PROVIDER, INGEST_MODE, STATUS_BACKEND,
  COLLECTION_METADATA_SOURCE, FACT_EVENT_EMITTER (+ supporting settings),
  validated in Settings.__post_init__ (fail-fast STATUS_BACKEND=local with
  INGEST_MODE=external); shared canonical.py.
- vector/payload_keys.py + acl_hash.py: cross-impl NAMESPACE/point_id (§2.2)
  and BLAKE2b-128 ACL hash (§11), pinned by fixtures shared with the
  document-processor repo.
- embedding/gateway_client.py: OpenAI-compatible GatewayProvider authenticating
  via M2M OIDC client-credentials (separate realm); manual-only registry entry.
- vector/collection_metadata.py: sentinel-point / API metadata source with env
  fallback.
- vector/queue/: hexagonal ingest producer ports + memory/NATS adapters
  (Postgres seam); INGEST_MODE=external publishes mcp.ingest.requested.{tenant}
  instead of the in-memory stream and skips the in-process processor pool. The
  lifespan becomes a composition root across both deployment branches.
- vector/queue/status.py: STATUS_BACKEND=bus subscriber feeding a StatusStore
  the vector-sync status endpoint reads.
- admin/payload_backfill.py: POST /api/v1/admin/payload-backfill (admin scope);
  processor writes the new payload keys; query-side ACL pre-filter gated behind
  ACL_PREFILTER_ENABLED (default off).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
2026-05-29 13:13:25 +02:00
co-authored by Claude Opus 4.8
parent c7da612f20
commit d883052fb8
37 changed files with 2534 additions and 55 deletions
@@ -0,0 +1,7 @@
"""Ingest-path ports & adapters (design §10, hexagonal)."""
from .factory import build_external_producer
from .memory import MemoryTaskProducer
from .ports import TaskProducer
__all__ = ["MemoryTaskProducer", "TaskProducer", "build_external_producer"]
@@ -0,0 +1,46 @@
"""Composition root for the ingest producer (design §10).
``build_external_producer`` is called by the lifespan only when
``INGEST_MODE=external``; local mode uses the in-memory stream directly (it
already satisfies :class:`TaskProducer`). The transport under ``external`` is
selected from the ``INGEST_BUS_URL`` scheme (``nats://`` now, ``postgres://``
later) so moving the external processor to Postgres needs no new INGEST_MODE.
"""
from __future__ import annotations
from urllib.parse import urlsplit
from ...config import Settings
from .ports import TaskProducer
def _transport_for(url: str) -> str:
scheme = urlsplit(url).scheme.lower()
if scheme.startswith("postgres"):
return "postgres"
return "nats"
async def build_external_producer(settings: Settings) -> TaskProducer:
"""Build the external-ingest producer for the configured transport.
Precondition: ``settings.ingest_mode == "external"`` (so __post_init__ has
guaranteed ``ingest_bus_url`` and ``tenant_id`` are set).
"""
assert settings.ingest_bus_url is not None
assert settings.tenant_id is not None
transport = _transport_for(settings.ingest_bus_url)
if transport == "postgres":
from .postgres import PostgresTaskProducer # noqa: PLC0415
return await PostgresTaskProducer.connect(settings)
from .nats import NatsTaskProducer # noqa: PLC0415
return await NatsTaskProducer.connect(
url=settings.ingest_bus_url,
tenant_id=settings.tenant_id,
num_replicas=settings.ingest_bus_num_replicas,
)
@@ -0,0 +1,45 @@
"""In-memory ``TaskProducer`` — the default (local) ingest transport.
A thin adapter over anyio's ``MemoryObjectSendStream`` so the local path is an
explicit :class:`TaskProducer` (rather than relying on structural typing of a
third-party class). Semantics are identical to using the stream directly:
``send`` enqueues, ``clone`` yields an independent per-user handle, ``async
with`` / ``aclose`` close the (cloned) send end so the processor pool's
receivers observe end-of-stream.
"""
from __future__ import annotations
from types import TracebackType
from typing import TYPE_CHECKING
from anyio.streams.memory import MemoryObjectSendStream
if TYPE_CHECKING:
from ..scanner import DocumentTask
class MemoryTaskProducer:
def __init__(self, stream: MemoryObjectSendStream[DocumentTask]):
self._stream = stream
async def send(self, task: DocumentTask, /) -> None:
await self._stream.send(task)
def clone(self) -> MemoryTaskProducer:
return MemoryTaskProducer(self._stream.clone())
async def __aenter__(self) -> MemoryTaskProducer:
await self._stream.__aenter__()
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
) -> None:
await self._stream.__aexit__(exc_type, exc, tb)
async def aclose(self) -> None:
await self._stream.aclose()
+147
View File
@@ -0,0 +1,147 @@
"""NATS JetStream ``TaskProducer`` — external ingest transport (design §3.4).
Publishes ``mcp.ingest.requested.{tenant_id}`` for the external
document-processor to consume. Translates the in-process ``DocumentTask`` into
the wire ``IngestMessage`` schema (mirrored in astrolabe-cloud-website's
``bus/messages.py``), with the JetStream ``Nats-Msg-Id`` dedup header per §3.4.
This server is only the *producer* on this transport; the document-processor
owns the consumer. ``nats-py`` is imported lazily so deployments that never
enable external ingest don't pay the import.
"""
from __future__ import annotations
import hashlib
import logging
from datetime import datetime, timezone
from types import TracebackType
from typing import TYPE_CHECKING, Any
from ...canonical import canonical_json
if TYPE_CHECKING:
from ..scanner import DocumentTask
logger = logging.getLogger(__name__)
STREAM_NAME = "mcp"
INGEST_SUBJECT_PREFIX = "mcp.ingest.requested"
def _modified_at_rfc3339(modified_at: int) -> str:
"""DocumentTask.modified_at is an epoch int (0 for deletes)."""
return datetime.fromtimestamp(int(modified_at), tz=timezone.utc).isoformat()
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)."""
return task.etag or str(task.modified_at)
def msg_id(tenant_id: str, doc_id: str, modified_at_rfc3339: str) -> str:
"""JetStream dedup header per §3.4. SHA-256 over canonical JSON (NOT the
BLAKE2b helper) — it is an opaque external header, not a stored field."""
return hashlib.sha256(
canonical_json(
{
"tenant_id": tenant_id,
"doc_id": doc_id,
"modified_at": modified_at_rfc3339,
}
)
).hexdigest()
class NatsTaskProducer:
"""Publishes ingest requests to NATS JetStream."""
def __init__(self, nc: Any, js: Any, tenant_id: str):
self._nc = nc
self._js = js
self.tenant_id = tenant_id
@classmethod
async def connect(
cls, *, url: str, tenant_id: str, num_replicas: int = 1
) -> NatsTaskProducer:
import nats # noqa: PLC0415 (lazy: optional dependency for external mode)
nc = await nats.connect(url)
js = nc.jetstream()
await cls._ensure_stream(js, num_replicas)
logger.info("Connected NATS ingest producer: url=%s, tenant=%s", url, tenant_id)
return cls(nc, js, tenant_id)
@staticmethod
async def _ensure_stream(js: Any, num_replicas: int) -> None:
# noqa: PLC0415 — nats.js types are only importable once nats-py is present.
from nats.js.api import RetentionPolicy, StreamConfig # noqa: PLC0415
config = StreamConfig(
name=STREAM_NAME,
subjects=["mcp.>"],
retention=RetentionPolicy.LIMITS,
num_replicas=num_replicas,
)
try:
await js.add_stream(config=config)
logger.info("nats.stream_created stream=%s", STREAM_NAME)
except Exception as exc:
# add_stream is idempotent in spirit but errors when the stream
# already exists; treat as benign (mirrors the processor's
# ensure_stream). A genuinely broken broker surfaces on publish.
logger.info("nats.stream_exists_or_unavailable detail=%s", exc)
def ingest_message(self, task: DocumentTask) -> dict[str, Any]:
"""DocumentTask → wire IngestMessage dict (mirrors the sibling schema)."""
return {
"tenant_id": self.tenant_id,
"doc_id": task.doc_id,
"content_hash": _content_hash(task),
"modified_at": _modified_at_rfc3339(task.modified_at),
"doc_type": task.doc_type,
"operation": task.operation,
"user_id": task.user_id,
"file_path": task.file_path,
}
async def send(self, task: DocumentTask) -> None:
message = self.ingest_message(task)
subject = f"{INGEST_SUBJECT_PREFIX}.{self.tenant_id}"
headers = {
"Nats-Msg-Id": msg_id(self.tenant_id, task.doc_id, message["modified_at"])
}
await self._js.publish(subject, canonical_json(message), headers=headers)
# The scanner/oauth_sync use the producer as a clone-able async context
# manager (memory-stream semantics). The bus connection is owned by the
# lifespan, so cloning shares it and __aexit__ is a no-op (close happens via
# aclose() on shutdown).
def clone(self) -> NatsTaskProducer:
return self
async def __aenter__(self) -> NatsTaskProducer:
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
) -> None:
return None
async def aclose(self) -> None:
# Per-handle close (e.g. a per-user scanner clone exiting). The bus
# connection is shared and owned by the lifespan, so this is a no-op;
# the connection is torn down once via ``drain()`` on shutdown.
return None
async def drain(self) -> None:
"""Drain + close the shared NATS connection (lifespan shutdown only)."""
try:
await self._nc.drain()
except Exception:
logger.warning("NATS drain on shutdown failed", exc_info=True)
@@ -0,0 +1,59 @@
"""Ingest-path ports (design §10, hexagonal).
A ``TaskProducer`` is where the scanner + webhook receiver send a
``DocumentTask``. The transport behind it is swappable:
- the in-process anyio ``MemoryObjectSendStream`` (local ingest — the default),
- ``NatsTaskProducer`` (external ingest → the document-processor), and
- a future Postgres-queue producer (seam only; the *external* processor owns the
consume side — see ``postgres.py``).
The protocol is exactly the surface the scanner/oauth_sync already use on the
memory stream (``send`` + ``clone`` + ``async with``), so both adapters drop in
with only a type-annotation change at the call sites. There is intentionally NO
consumer port: the MCP server's only in-process consumer is the memory stream;
when ingest is external the document-processor is the consumer, not this server.
"""
from __future__ import annotations
from types import TracebackType
from typing import TYPE_CHECKING, Protocol, runtime_checkable
if TYPE_CHECKING:
from ..scanner import DocumentTask
@runtime_checkable
class TaskProducer(Protocol):
"""Sink for scanner/webhook ``DocumentTask``s (see module docstring)."""
# Positional-only so anyio's MemoryObjectSendStream.send(item) structurally
# satisfies this protocol (its parameter is named "item", not "task").
async def send(self, task: DocumentTask, /) -> None: ...
def clone(self) -> TaskProducer:
"""Return a producer handle for one user's scanner (multi-user mode).
For the memory stream this is a real clone (each closed independently);
for the bus it returns ``self`` (one shared connection).
"""
...
async def __aenter__(self) -> TaskProducer: ...
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
) -> None: ...
async def aclose(self) -> None:
"""Close *this* handle (e.g. a per-user clone when its scanner exits).
For the memory stream this closes the clone; for the shared bus
connection it is a no-op (the connection is owned by the lifespan,
which drains it once on shutdown).
"""
...
@@ -0,0 +1,49 @@
"""Postgres-queue ``TaskProducer`` — documented seam, not implemented.
The external document-processor may later drain a Postgres-backed queue instead
of NATS to limit NATS operational overhead. Processing stays *external*; only
the transport changes — so on this server it would be a drop-in producer swap.
The consume side + the queue-table migration belong to that processor-side
refactor (cross-repo), NOT here. This stub exists so the transport value and the
``TaskProducer`` Protocol conformance are testable today.
"""
from __future__ import annotations
from types import TracebackType
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from ..scanner import DocumentTask
_NOT_IMPLEMENTED = (
"Postgres ingest transport is a documented seam. The external "
"document-processor owns the Postgres-drain refactor (transport swap only; "
"processing stays external). Use INGEST_BUS_URL=nats://… for now."
)
class PostgresTaskProducer:
@classmethod
async def connect(cls, settings: Any) -> PostgresTaskProducer:
raise NotImplementedError(_NOT_IMPLEMENTED)
async def send(self, task: DocumentTask) -> None: # pragma: no cover
raise NotImplementedError(_NOT_IMPLEMENTED)
def clone(self) -> PostgresTaskProducer: # pragma: no cover
return self
async def __aenter__(self) -> PostgresTaskProducer: # pragma: no cover
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
) -> None: # pragma: no cover
return None
async def aclose(self) -> None: # pragma: no cover
return None
+163
View File
@@ -0,0 +1,163 @@
"""Status surface for ingest jobs (design §10.1, ``STATUS_BACKEND``).
- ``local``: in-process job state — the memory-stream buffer (today's behavior,
read directly by the status endpoint).
- ``bus``: a background subscriber consumes
``mcp.document.{ready,failed,reparsed}.{tenant_id}`` into a bounded in-process
:class:`StatusStore` that the status endpoint / ``nc_get_vector_sync_status``
read.
**Honest constraint (design §10.2 / decision):** MCP progress notifications
(``ctx.report_progress``) can only be emitted inside an *active tool-call
request*; a background subscriber has no ``ctx`` and the MCP SDK exposes no
out-of-band push. So "surface events as MCP progress notifications" is delivered
via this store (polled by the status endpoint / a tool), not an unsolicited
server push. True server-initiated progress / SSE is a follow-up — the
``on_event`` callback seam is left in place for it.
"""
from __future__ import annotations
import logging
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Callable
if TYPE_CHECKING:
import anyio
from anyio.abc import TaskStatus
logger = logging.getLogger(__name__)
# Terminal/intermediate document states carried on mcp.document.* subjects.
_VALID_STATES = {"ready", "failed", "reparsed"}
class StatusStore:
"""Bounded LRU of recent document states keyed by ``doc_id``."""
def __init__(self, max_size: int = 10_000):
self._entries: OrderedDict[str, dict[str, Any]] = OrderedDict()
self._max = max_size
def record(
self,
doc_id: str,
state: str,
*,
content_hash: str | None = None,
transitioned_at: str | None = None,
) -> None:
self._entries[doc_id] = {
"state": state,
"content_hash": content_hash,
"transitioned_at": transitioned_at,
}
self._entries.move_to_end(doc_id)
while len(self._entries) > self._max:
self._entries.popitem(last=False)
def get(self, doc_id: str) -> dict[str, Any] | None:
return self._entries.get(doc_id)
def counts(self) -> dict[str, int]:
out: dict[str, int] = {}
for entry in self._entries.values():
out[entry["state"]] = out.get(entry["state"], 0) + 1
return out
def __len__(self) -> int:
return len(self._entries)
def state_from_subject(subject: str) -> str | None:
"""``mcp.document.<state>.<tenant_id>`` → ``<state>`` (or None if unknown)."""
parts = subject.split(".")
if len(parts) >= 4 and parts[0] == "mcp" and parts[1] == "document":
state = parts[2]
if state in _VALID_STATES:
return state
return None
class NatsStatusSubscriber:
"""Consumes ``mcp.document.*.{tenant_id}`` into a :class:`StatusStore`."""
def __init__(
self,
nc: Any,
js: Any,
tenant_id: str,
store: StatusStore,
on_event: Callable[[str, str], None] | None = None,
):
self._nc = nc
self._js = js
self.tenant_id = tenant_id
self.store = store
# on_event(doc_id, state) — seam for a future SSE / progress bridge.
self._on_event = on_event
def handle_message(self, subject: str, data: bytes) -> None:
"""Parse one status message into the store. Unit-testable without NATS."""
import json # noqa: PLC0415
state = state_from_subject(subject)
if state is None:
logger.warning("status.unknown_subject subject=%s", subject)
return
try:
payload = json.loads(data)
doc_id = payload["doc_id"]
except Exception:
logger.warning("status.bad_message subject=%s", subject, exc_info=True)
return
self.store.record(
doc_id,
state,
content_hash=payload.get("content_hash"),
transitioned_at=payload.get("transitioned_at"),
)
if self._on_event is not None:
self._on_event(doc_id, state)
@classmethod
async def connect(
cls, *, url: str, tenant_id: str, store: StatusStore
) -> NatsStatusSubscriber:
import nats # noqa: PLC0415
nc = await nats.connect(url)
js = nc.jetstream()
return cls(nc, js, tenant_id, store)
async def run(
self,
shutdown_event: anyio.Event,
*,
task_status: TaskStatus = None, # type: ignore[assignment]
) -> None:
"""Durable pull-consumer loop. Requires a live broker (integration)."""
import anyio # noqa: PLC0415
subject = f"mcp.document.*.{self.tenant_id}"
sub = await self._js.pull_subscribe(
subject, durable=f"mcp-status-{self.tenant_id}"
)
if task_status is not None:
task_status.started()
while not shutdown_event.is_set():
try:
msgs = await sub.fetch(batch=16, timeout=5)
except Exception:
# fetch timeout when idle — loop and re-check shutdown.
await anyio.sleep(0)
continue
for msg in msgs:
self.handle_message(msg.subject, msg.data)
await msg.ack()
async def aclose(self) -> None:
try:
await self._nc.drain()
except Exception:
logger.warning("NATS status subscriber drain failed", exc_info=True)