fix: address PR #836 review — forward task_producer to MCP contexts + cleanups

🔴 nc_get_vector_sync_status reported pending=0 for INGEST_QUEUE=postgres: the
AppContext/OAuthAppContext per-session yields snapshotted the stream fields but
never forwarded task_producer, so lifespan_ctx.task_producer was always None.
Convert task_producer to a @property that reads _vector_sync_state live (like
eviction_task_group), removing the snapshot field so the yields can't drop it.
Add a regression test pinning the contract on both contexts.

🟡 Remove the unused _RECLAIM_TASK_NAME constant.
🟡 get_procrastinate_conninfo: warn + document that DATABASE_URL query params
   (application_name, connect_timeout, …) are dropped.
🟡 worker: open the procrastinate App once — apply_ingest_queue_schema gains
   manage_connection=False so the worker reuses its own open connector instead
   of a redundant open/close before run_worker_async.

🟢 Clarify the apply-schema broad-except comment (non-race errors re-raise) and
   document the deliberate Any typing in ingest_status.get_ingest_pending.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
2026-06-03 12:40:50 +02:00
co-authored by Claude Opus 4.8
parent 63e073c224
commit cfdef3c2c5
6 changed files with 110 additions and 26 deletions
+18 -6
View File
@@ -355,10 +355,18 @@ class AppContext:
storage: "RefreshTokenStorage | None" = None storage: "RefreshTokenStorage | None" = None
document_send_stream: MemoryObjectSendStream | None = None document_send_stream: MemoryObjectSendStream | None = None
document_receive_stream: MemoryObjectReceiveStream | None = None document_receive_stream: MemoryObjectReceiveStream | None = None
task_producer: "TaskProducer | None" = None
shutdown_event: anyio.Event | None = None shutdown_event: anyio.Event | None = None
scanner_wake_event: anyio.Event | None = None scanner_wake_event: anyio.Event | None = None
@property
def task_producer(self) -> "TaskProducer | None":
# Read dynamically from the module-level singleton (like
# eviction_task_group) rather than snapshotting at yield time — that way
# a session can't observe a stale ``None`` and the per-session yields
# can't forget to forward it (the bug this property replaces). The
# vector-sync status tool reads this for postgres-backend job counts.
return _vector_sync_state.task_producer
@property @property
def eviction_task_group(self) -> TaskGroup | None: def eviction_task_group(self) -> TaskGroup | None:
# Read dynamically from the module-level singleton instead of # Read dynamically from the module-level singleton instead of
@@ -381,10 +389,14 @@ class OAuthAppContext:
server_client_id: str | None = None # MCP server's OAuth client ID (static or DCR) server_client_id: str | None = None # MCP server's OAuth client ID (static or DCR)
document_send_stream: MemoryObjectSendStream | None = None document_send_stream: MemoryObjectSendStream | None = None
document_receive_stream: MemoryObjectReceiveStream | None = None document_receive_stream: MemoryObjectReceiveStream | None = None
task_producer: "TaskProducer | None" = None
shutdown_event: anyio.Event | None = None shutdown_event: anyio.Event | None = None
scanner_wake_event: anyio.Event | None = None scanner_wake_event: anyio.Event | None = None
@property
def task_producer(self) -> "TaskProducer | None":
# See AppContext.task_producer for rationale.
return _vector_sync_state.task_producer
@property @property
def eviction_task_group(self) -> TaskGroup | None: def eviction_task_group(self) -> TaskGroup | None:
# See AppContext.eviction_task_group for rationale. # See AppContext.eviction_task_group for rationale.
@@ -606,8 +618,8 @@ async def app_lifespan_basic(server: FastMCP) -> AsyncIterator[AppContext]:
document_receive_stream=_vector_sync_state.document_receive_stream, document_receive_stream=_vector_sync_state.document_receive_stream,
shutdown_event=_vector_sync_state.shutdown_event, shutdown_event=_vector_sync_state.shutdown_event,
scanner_wake_event=_vector_sync_state.scanner_wake_event, scanner_wake_event=_vector_sync_state.scanner_wake_event,
# eviction_task_group is exposed via @property (reads # task_producer and eviction_task_group are exposed via @property
# _vector_sync_state at access time, not snapshot). # (read _vector_sync_state at access time, not snapshot).
) )
finally: finally:
logger.info("Shutting down BasicAuth session") logger.info("Shutting down BasicAuth session")
@@ -1239,8 +1251,8 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
document_receive_stream=_vector_sync_state.document_receive_stream, document_receive_stream=_vector_sync_state.document_receive_stream,
shutdown_event=_vector_sync_state.shutdown_event, shutdown_event=_vector_sync_state.shutdown_event,
scanner_wake_event=_vector_sync_state.scanner_wake_event, scanner_wake_event=_vector_sync_state.scanner_wake_event,
# eviction_task_group is exposed via @property (reads # task_producer and eviction_task_group are exposed via
# _vector_sync_state at access time, not snapshot). # @property (read _vector_sync_state at access time).
) )
finally: finally:
logger.info("Shutting down MCP server") logger.info("Shutting down MCP server")
+5 -2
View File
@@ -326,9 +326,12 @@ def worker(concurrency: int | None):
app = get_procrastinate_app() app = get_procrastinate_app()
async def _run() -> None: async def _run() -> None:
# Defensive apply (the always-on API pod is the authoritative applier). # Open the connector pool once and reuse it for both the defensive
await apply_ingest_queue_schema(app) # schema apply (the always-on API pod is the authoritative applier) and
# the worker loop — manage_connection=False avoids a redundant
# open/close cycle on startup.
async with app.open_async(): async with app.open_async():
await apply_ingest_queue_schema(app, manage_connection=False)
click.echo( click.echo(
f"Ingest worker started: queue={INGEST_QUEUE_NAME} concurrency={workers}" f"Ingest worker started: queue={INGEST_QUEUE_NAME} concurrency={workers}"
) )
+14
View File
@@ -1437,6 +1437,13 @@ def get_procrastinate_conninfo(database_url: str | None = None) -> str:
(set ``DATABASE_URL`` to a ``+psycopg`` URL) rather than rewriting the driver (set ``DATABASE_URL`` to a ``+psycopg`` URL) rather than rewriting the driver
in code — see charts repo, not this repo. in code — see charts repo, not this repo.
Only the host/port/dbname/user/password components are forwarded; any
``?key=value`` query parameters on the URL (e.g. ``application_name``,
``connect_timeout``) are **dropped** — TLS is set separately via
:func:`_pg_ssl_params`, and SQLAlchemy-specific query options don't map
cleanly to libpq keywords. A warning is logged when params are dropped so
operators aren't surprised.
Raises ``ValueError`` for a non-Postgres URL — procrastinate is Postgres-only. Raises ``ValueError`` for a non-Postgres URL — procrastinate is Postgres-only.
""" """
from psycopg.conninfo import make_conninfo # noqa: PLC0415 from psycopg.conninfo import make_conninfo # noqa: PLC0415
@@ -1449,6 +1456,13 @@ def get_procrastinate_conninfo(database_url: str | None = None) -> str:
f"got driver {url.drivername!r}" f"got driver {url.drivername!r}"
) )
if url.query:
logging.getLogger(__name__).warning(
"Dropping DATABASE_URL query parameters not forwarded to the "
"procrastinate connector: %s",
", ".join(sorted(url.query)),
)
params: dict[str, str] = {} params: dict[str, str] = {}
if url.host: if url.host:
params["host"] = url.host params["host"] = url.host
@@ -36,6 +36,12 @@ async def get_ingest_pending(
) -> IngestPending: ) -> IngestPending:
"""Compute outstanding ingest work for the configured queue backend. """Compute outstanding ingest work for the configured queue backend.
``task_producer`` and ``document_receive_stream`` are intentionally typed
``Any``: they're duck-typed across backends. Only ``ProcrastinateTaskProducer``
exposes ``job_counts`` (the ``TaskProducer`` protocol doesn't), and the memory
backend reads the anyio stream's ``statistics()`` — so no single concrete type
or Protocol fits both branches, and we probe with ``hasattr`` instead.
Never raises — a status surface must stay available even if the queue is Never raises — a status surface must stay available even if the queue is
unreachable; failures degrade to ``pending=0``. unreachable; failures degrade to ``pending=0``.
""" """
@@ -54,7 +54,6 @@ INGEST_QUEUE_NAME = "ingest"
# Blueprint namespace → registered task names are prefixed ``ingest:``. # Blueprint namespace → registered task names are prefixed ``ingest:``.
_NAMESPACE = "ingest" _NAMESPACE = "ingest"
INGEST_TASK_NAME = f"{_NAMESPACE}:process_document" INGEST_TASK_NAME = f"{_NAMESPACE}:process_document"
_RECLAIM_TASK_NAME = f"{_NAMESPACE}:reclaim_stalled_jobs"
# A crashed worker leaves its job in ``doing``; reclaim it once its (per-worker) # A crashed worker leaves its job in ``doing``; reclaim it once its (per-worker)
# heartbeat is this many seconds stale. Sized well above the longest expected # heartbeat is this many seconds stale. Sized well above the longest expected
@@ -217,8 +216,8 @@ async def _ingest_schema_present(app: App) -> bool:
return bool(row["present"]) return bool(row["present"])
async def apply_ingest_queue_schema(app: App | None = None) -> None: async def _apply_ingest_queue_schema_open(app: App) -> None:
"""Create procrastinate's tables on a fresh database (apply-if-absent). """Apply the ingest-queue schema on an already-open connector (apply-if-absent).
procrastinate's ``schema.sql`` uses bare ``CREATE TYPE``/``CREATE TABLE`` procrastinate's ``schema.sql`` uses bare ``CREATE TYPE``/``CREATE TABLE``
(not ``IF NOT EXISTS``), so it errors if re-applied — it is meant to run (not ``IF NOT EXISTS``), so it errors if re-applied — it is meant to run
@@ -230,25 +229,41 @@ async def apply_ingest_queue_schema(app: App | None = None) -> None:
lock: Postgres DDL is transactional and procrastinate applies the whole lock: Postgres DDL is transactional and procrastinate applies the whole
schema in one transaction, so a pod that loses the race rolls back cleanly schema in one transaction, so a pod that loses the race rolls back cleanly
and we treat the resulting error as benign once the schema is present. and we treat the resulting error as benign once the schema is present.
"""
if await _ingest_schema_present(app):
logger.debug("ingest queue schema already present; skipping apply")
return
try:
await app.schema_manager.apply_schema_async()
logger.info("Applied procrastinate ingest queue schema")
except Exception:
# The apply runs in a single transaction, so any failure rolls back
# atomically (no partial schema). The only benign case is losing the
# create race to another pod — confirmed by re-checking presence. Any
# other failure (network, auth, …) leaves the schema absent, so this
# branch re-raises it rather than masking it.
if await _ingest_schema_present(app):
logger.info("Ingest queue schema applied concurrently by another pod")
return
raise
Opens a short-lived connection, so it is safe to call from the CLI
(``db upgrade`` / worker startup). async def apply_ingest_queue_schema(
app: App | None = None, *, manage_connection: bool = True
) -> None:
"""Create procrastinate's tables on a fresh database (apply-if-absent).
By default opens a short-lived connection, so it is safe to call standalone
from the CLI ``db upgrade`` path. Pass ``manage_connection=False`` when the
caller already holds an open connector (the ``worker`` command opens the App
once and reuses it) to avoid a redundant open/close cycle.
""" """
app = app or get_procrastinate_app() app = app or get_procrastinate_app()
if not manage_connection:
await _apply_ingest_queue_schema_open(app)
return
async with app.open_async(): async with app.open_async():
if await _ingest_schema_present(app): await _apply_ingest_queue_schema_open(app)
logger.debug("ingest queue schema already present; skipping apply")
return
try:
await app.schema_manager.apply_schema_async()
logger.info("Applied procrastinate ingest queue schema")
except Exception:
# A racing pod likely committed the schema while our transaction
# rolled back atomically. Benign iff the schema is now present.
if await _ingest_schema_present(app):
logger.info("Ingest queue schema applied concurrently by another pod")
return
raise
# Job-status keys procrastinate flattens into each list_queues row (alongside # Job-status keys procrastinate flattens into each list_queues row (alongside
@@ -0,0 +1,34 @@
"""Regression test for the lifespan context `task_producer` exposure (Deck #183).
`nc_get_vector_sync_status` reads `lifespan_ctx.task_producer` for postgres-backend
job counts. It was previously a snapshot dataclass field the per-session yields
forgot to populate, so the tool always reported `pending=0` on the postgres
backend. It is now a `@property` that reads the module singleton live (like
`eviction_task_group`); these tests pin that contract.
"""
from typing import cast
import pytest
import nextcloud_mcp_server.app as app_module
from nextcloud_mcp_server.app import AppContext, OAuthAppContext
from nextcloud_mcp_server.client import NextcloudClient
pytestmark = pytest.mark.unit
def test_app_context_task_producer_reads_vector_sync_state(monkeypatch):
sentinel = object()
monkeypatch.setattr(app_module._vector_sync_state, "task_producer", sentinel)
ctx = AppContext(client=cast(NextcloudClient, None))
assert ctx.task_producer is sentinel
def test_oauth_app_context_task_producer_reads_vector_sync_state(monkeypatch):
sentinel = object()
monkeypatch.setattr(app_module._vector_sync_state, "task_producer", sentinel)
ctx = OAuthAppContext(
nextcloud_host="https://example.test", token_verifier=object()
)
assert ctx.task_producer is sentinel