refactor: migrate vector sync from asyncio.Queue to anyio memory object streams

Replace asyncio.Queue with anyio.create_memory_object_stream() throughout
the vector sync system for better library consistency and improved shutdown
semantics.

## Changes Made

**scanner.py**:
- Changed parameter type from `asyncio.Queue` to `MemoryObjectSendStream[DocumentTask]`
- Replaced all `await document_queue.put()` calls with `await send_stream.send()`
- Wrapped scanner loop in `async with send_stream:` context manager for automatic cleanup
- Updated log messages: "Queued" → "Sent"
- Removed `import asyncio` (no longer needed)

**processor.py**:
- Changed parameter type from `asyncio.Queue` to `MemoryObjectReceiveStream[DocumentTask]`
- Replaced `asyncio.wait_for(document_queue.get(), timeout=1.0)` with `anyio.fail_after(1.0)` + `await receive_stream.receive()`
- Removed all `document_queue.task_done()` calls (not needed with streams)
- Added `anyio.EndOfStream` exception handling for graceful shutdown when scanner closes
- Removed `import asyncio` (no longer needed)

**app.py**:
- Removed `import asyncio` from top-level imports
- Added `from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream`
- Updated AppContext dataclass:
  - Replaced `document_queue: Optional[asyncio.Queue]` with:
    - `document_send_stream: Optional[MemoryObjectSendStream]`
    - `document_receive_stream: Optional[MemoryObjectReceiveStream]`
- Updated `app_lifespan_basic()`:
  - Replaced `asyncio.Queue(maxsize=...)` with `anyio.create_memory_object_stream(max_buffer_size=...)`
  - Pass `send_stream` to scanner_task
  - Pass `receive_stream.clone()` to each processor_task (enables multiple consumers)
  - Updated AppContext yield to include both streams
- Updated `starlette_lifespan()`:
  - Same changes as app_lifespan_basic for streamable-http transport
  - Removed `import asyncio as asyncio_module` (no longer needed)
  - Updated app.state storage to use send_stream and receive_stream

**semantic.py**:
- Updated `nc_get_vector_sync_status()` tool:
  - Access `document_receive_stream` instead of `document_queue` from lifespan context
  - Use `stream_stats.current_buffer_used` instead of `queue.qsize()` for pending count
  - More reliable metrics (qsize() was not guaranteed accurate)

## Benefits

1. **Library Consistency**: Pure anyio throughout codebase (was mixing asyncio.Queue with anyio.Event and anyio.create_task_group)
2. **Graceful Shutdown**: `async with send_stream:` automatically closes stream on exit, signaling EndOfStream to all processors
3. **Better Timeout Handling**: `anyio.fail_after()` is more idiomatic than `asyncio.wait_for()`
4. **Stream Cloning**: Easy to add multiple consumers via `receive_stream.clone()`
5. **Better Statistics**: `.statistics()` provides accurate buffer metrics (qsize() was unreliable)
6. **Type Safety**: Separate send/receive types prevent accidental misuse
7. **No task_done() tracking**: Streams handle completion automatically

## Testing

-  All 69 unit tests passing
-  All 5 smoke tests passing
-  No regressions in functionality
-  Graceful shutdown behavior improved

## References

- https://anyio.readthedocs.io/en/stable/why.html#queue-fix
- https://anyio.readthedocs.io/en/stable/streams.html#memory-object-streams

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
2025-11-09 06:43:44 +01:00
co-authored by Claude
parent 4b026e9aa0
commit 72232f937a
4 changed files with 83 additions and 78 deletions
+20 -16
View File
@@ -1,4 +1,3 @@
import asyncio
import logging
import os
from collections.abc import AsyncIterator
@@ -13,6 +12,7 @@ import anyio
import click
import httpx
import uvicorn
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.server.auth.settings import AuthSettings
from mcp.server.fastmcp import Context, FastMCP
from pydantic import AnyHttpUrl
@@ -211,7 +211,8 @@ class AppContext:
"""Application context for BasicAuth mode."""
client: NextcloudClient
document_queue: Optional[asyncio.Queue] = None
document_send_stream: Optional[MemoryObjectSendStream] = None
document_receive_stream: Optional[MemoryObjectReceiveStream] = None
shutdown_event: Optional[anyio.Event] = None
scanner_wake_event: Optional[anyio.Event] = None
@@ -404,7 +405,9 @@ async def app_lifespan_basic(server: FastMCP) -> AsyncIterator[AppContext]:
)
# Initialize shared state
document_queue = asyncio.Queue(maxsize=settings.vector_sync_queue_max_size)
send_stream, receive_stream = anyio.create_memory_object_stream(
max_buffer_size=settings.vector_sync_queue_max_size
)
shutdown_event = anyio.Event()
scanner_wake_event = anyio.Event()
@@ -413,19 +416,19 @@ async def app_lifespan_basic(server: FastMCP) -> AsyncIterator[AppContext]:
# Start scanner task
tg.start_soon(
scanner_task,
document_queue,
send_stream,
shutdown_event,
scanner_wake_event,
client,
username,
)
# Start processor pool
# Start processor pool (each gets a cloned receive stream)
for i in range(settings.vector_sync_processor_workers):
tg.start_soon(
processor_task,
i,
document_queue,
receive_stream.clone(),
shutdown_event,
client,
username,
@@ -439,7 +442,8 @@ async def app_lifespan_basic(server: FastMCP) -> AsyncIterator[AppContext]:
try:
yield AppContext(
client=client,
document_queue=document_queue,
document_send_stream=send_stream,
document_receive_stream=receive_stream,
shutdown_event=shutdown_event,
scanner_wake_event=scanner_wake_event,
)
@@ -1009,8 +1013,6 @@ def get_app(transport: str = "sse", enabled_apps: list[str] | None = None):
# Start background vector sync tasks for BasicAuth mode (ADR-007)
# For streamable-http transport, FastMCP lifespan isn't automatically triggered
# so we manually start background tasks here if vector sync is enabled
import asyncio as asyncio_module
import anyio as anyio_module
settings = get_settings()
@@ -1029,21 +1031,23 @@ def get_app(transport: str = "sse", enabled_apps: list[str] | None = None):
client = NextcloudClient.from_env()
# Initialize shared state
document_queue = asyncio_module.Queue(
maxsize=settings.vector_sync_queue_max_size
send_stream, receive_stream = anyio_module.create_memory_object_stream(
max_buffer_size=settings.vector_sync_queue_max_size
)
shutdown_event = anyio_module.Event()
scanner_wake_event = anyio_module.Event()
# Store in app state for access from routes (ADR-007)
app.state.document_queue = document_queue
app.state.document_send_stream = send_stream
app.state.document_receive_stream = receive_stream
app.state.shutdown_event = shutdown_event
app.state.scanner_wake_event = scanner_wake_event
# Also share with browser_app for /user/page route
for route in app.routes:
if isinstance(route, Mount) and route.path == "/user":
route.app.state.document_queue = document_queue
route.app.state.document_send_stream = send_stream
route.app.state.document_receive_stream = receive_stream
route.app.state.shutdown_event = shutdown_event
route.app.state.scanner_wake_event = scanner_wake_event
logger.info(
@@ -1056,19 +1060,19 @@ def get_app(transport: str = "sse", enabled_apps: list[str] | None = None):
# Start scanner task
tg.start_soon(
scanner_task,
document_queue,
send_stream,
shutdown_event,
scanner_wake_event,
client,
username,
)
# Start processor pool
# Start processor pool (each gets a cloned receive stream)
for i in range(settings.vector_sync_processor_workers):
tg.start_soon(
processor_task,
i,
document_queue,
receive_stream.clone(),
shutdown_event,
client,
username,