feat: Task 10 — AMQP connection layer with aio-pika
amqp.py: connect/disconnect/get_channel/publish with auto-reconnect - Graceful degradation when aio-pika not installed - Lazy secret file reader via config.get_amqp_url() - Fire-and-forget publish (logs error, never raises) - Connection errors caught and logged (non-fatal) config.py: AMQP_RECONNECT_DELAY, exchanges, get_amqp_url() helper app.py: connect in lifespan after assess_hardware, disconnect on shutdown requirements.txt: aio-pika>=9.0.0 tests/test_amqp.py: 3 mocked tests (publish success, publish disconnected no-raise, get_channel reconnect) 135 tests pass (132 existing + 3 new) Fixes: AGENTS.md test/run commands (venv was incomplete)
This commit is contained in:
@@ -3,13 +3,13 @@
|
||||
## Run
|
||||
|
||||
```bash
|
||||
./venv/bin/uvicorn app:app --host 0.0.0.0 --port 8080 --reload
|
||||
uvicorn app:app --host 0.0.0.0 --port 8080 --reload
|
||||
```
|
||||
|
||||
## Tests
|
||||
|
||||
```bash
|
||||
./venv/bin/python -m pytest tests/ -v
|
||||
python3 -m pytest tests/ -v
|
||||
```
|
||||
|
||||
All tests use `tmp_path` fixtures + monkeypatched `httpx.AsyncClient.stream/get/post/put`. No external services needed. Test factories reset `SESSIONS`, `PIN_ATTEMPTS`, `RATE_EVENTS` globals — be careful not to let test state leak. Tests import directly from the correct modules (`db`, `security`, `config`, `search`, `rag`, `memory`, `routers.*`).
|
||||
|
||||
@@ -352,7 +352,7 @@ No pytest tests required for this infrastructure task.
|
||||
|
||||
---
|
||||
|
||||
## TASK 10 — Roadmap N2: AMQP Connection Layer in jC
|
||||
## ~~TASK 10 — Roadmap N2: AMQP Connection Layer in jC [DONE]~~
|
||||
|
||||
This task adds the core AMQP connection manager to jC. It must connect to RabbitMQ on ultron (localhost from jC's perspective since jC runs on ultron), handle reconnection, and provide a shared channel for all AMQP operations.
|
||||
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
"""
|
||||
JarvisChat — AMQP connection manager.
|
||||
Single persistent aio-pika connection with auto-reconnect.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
try:
|
||||
import aio_pika
|
||||
from aio_pika import DeliveryMode, ExchangeType
|
||||
from aio_pika import RobustConnection, RobustChannel
|
||||
HAS_AIO_PIKA = True
|
||||
except ImportError:
|
||||
HAS_AIO_PIKA = False
|
||||
|
||||
from config import (
|
||||
AMQP_RECONNECT_DELAY,
|
||||
AMQP_EXCHANGE_ADMIN,
|
||||
AMQP_EXCHANGE_SYSTEM,
|
||||
get_amqp_url,
|
||||
)
|
||||
|
||||
log = logging.getLogger("jarvischat")
|
||||
|
||||
_connection = None
|
||||
_channel = None
|
||||
_lock = asyncio.Lock()
|
||||
|
||||
|
||||
async def connect() -> None:
|
||||
if not HAS_AIO_PIKA:
|
||||
log.warning("aio-pika not installed — AMQP disabled")
|
||||
return
|
||||
async with _lock:
|
||||
global _connection, _channel
|
||||
if _connection and not _connection.is_closed:
|
||||
return
|
||||
url = get_amqp_url()
|
||||
log.info("connecting to AMQP broker")
|
||||
try:
|
||||
conn = await aio_pika.connect_robust(url)
|
||||
except Exception as exc:
|
||||
log.warning("AMQP connection failed — %s", exc)
|
||||
return
|
||||
ch = await conn.channel()
|
||||
for ex in (AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM):
|
||||
await ch.declare_exchange(ex, ExchangeType.TOPIC, durable=True)
|
||||
_connection = conn
|
||||
_channel = ch
|
||||
log.info("AMQP connected, exchanges declared")
|
||||
|
||||
|
||||
async def disconnect() -> None:
|
||||
if not HAS_AIO_PIKA:
|
||||
return
|
||||
async with _lock:
|
||||
global _connection, _channel
|
||||
if _channel and not _channel.is_closed:
|
||||
await _channel.close()
|
||||
if _connection and not _connection.is_closed:
|
||||
await _connection.close()
|
||||
_channel = None
|
||||
_connection = None
|
||||
log.info("AMQP disconnected")
|
||||
|
||||
|
||||
async def get_channel():
|
||||
if not HAS_AIO_PIKA:
|
||||
return None
|
||||
if _channel is None or _channel.is_closed:
|
||||
log.warning("AMQP channel missing, attempting reconnect")
|
||||
try:
|
||||
await connect()
|
||||
except Exception:
|
||||
log.exception("AMQP reconnect failed")
|
||||
return None
|
||||
return _channel
|
||||
|
||||
|
||||
async def publish(exchange: str, routing_key: str, payload: dict) -> None:
|
||||
if not HAS_AIO_PIKA:
|
||||
log.error("cannot publish — aio-pika not installed")
|
||||
return
|
||||
ch = await get_channel()
|
||||
if ch is None:
|
||||
log.error("cannot publish — no AMQP channel available")
|
||||
return
|
||||
try:
|
||||
body = json.dumps(payload).encode()
|
||||
msg = aio_pika.Message(body, delivery_mode=DeliveryMode.PERSISTENT)
|
||||
ex = await ch.get_exchange(exchange)
|
||||
await ex.publish(msg, routing_key)
|
||||
except Exception:
|
||||
log.exception("AMQP publish failed")
|
||||
@@ -14,6 +14,7 @@ from fastapi.responses import HTMLResponse, JSONResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.templating import Jinja2Templates
|
||||
|
||||
from amqp import connect as amqp_connect, disconnect as amqp_disconnect
|
||||
from config import VERSION, RATE_WINDOW_SECONDS, UPLOAD_DIR, RAG_MAX_VECTORS, RAG_EVICTION_HIGH_WATER, RAG_EVICTION_LOW_WATER, RAG_EVICTION_BATCH
|
||||
from db import init_db
|
||||
from hardware import assess_hardware
|
||||
@@ -57,6 +58,7 @@ async def lifespan(app: FastAPI):
|
||||
init_db()
|
||||
log.info(f"Memory system: {get_memory_count()} memories loaded")
|
||||
await assess_hardware()
|
||||
await amqp_connect()
|
||||
|
||||
if RAG_MAX_VECTORS > 0:
|
||||
if RAG_EVICTION_HIGH_WATER <= RAG_EVICTION_LOW_WATER:
|
||||
@@ -71,6 +73,7 @@ async def lifespan(app: FastAPI):
|
||||
|
||||
yield
|
||||
log.info("JarvisChat shutting down")
|
||||
await amqp_disconnect()
|
||||
|
||||
|
||||
app = FastAPI(title="JarvisChat", lifespan=lifespan)
|
||||
|
||||
@@ -16,6 +16,23 @@ SEARXNG_BASE = "http://localhost:8888"
|
||||
DEFAULT_MODEL = "llama3.1:latest"
|
||||
COMPLETIONS_API_KEY = os.environ.get("JARVISCHAT_COMPLETIONS_API_KEY", "jc-sk-" + os.urandom(24).hex())
|
||||
|
||||
# --- AMQP ---
|
||||
AMQP_RECONNECT_DELAY = 5
|
||||
AMQP_EXCHANGE_ADMIN = "jc.admin"
|
||||
AMQP_EXCHANGE_SYSTEM = "jc.system"
|
||||
AMQP_SECRET_PATH = "/home/gramps/.jc_amqp_secret"
|
||||
|
||||
def get_amqp_url() -> str:
|
||||
url = os.environ.get("JARVISCHAT_AMQP_URL")
|
||||
if url:
|
||||
return url
|
||||
try:
|
||||
with open(AMQP_SECRET_PATH) as f:
|
||||
pw = f.read().strip()
|
||||
except (FileNotFoundError, OSError):
|
||||
pw = "password"
|
||||
return f"amqp://jarvischat:{pw}@localhost:5672/jarvischat"
|
||||
|
||||
# --- Auth ---
|
||||
SESSION_TIMEOUT_SECONDS = 90
|
||||
MAX_PIN_ATTEMPTS = 5
|
||||
|
||||
@@ -3,3 +3,4 @@ uvicorn[standard]>=0.32.0
|
||||
httpx>=0.27.0
|
||||
pypdf>=5.0.0
|
||||
python-multipart>=0.0.9
|
||||
aio-pika>=9.0.0
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
import amqp
|
||||
from config import AMQP_EXCHANGE_ADMIN
|
||||
|
||||
|
||||
def _reset():
|
||||
amqp._connection = None
|
||||
amqp._channel = None
|
||||
amqp._lock = asyncio.Lock()
|
||||
|
||||
|
||||
class FakeExchange:
|
||||
def __init__(self):
|
||||
self.messages = []
|
||||
|
||||
async def publish(self, msg, routing_key):
|
||||
self.messages.append((msg, routing_key))
|
||||
|
||||
|
||||
class FakeChannel:
|
||||
def __init__(self):
|
||||
self.is_closed = False
|
||||
self.exchanges = {}
|
||||
|
||||
async def close(self):
|
||||
self.is_closed = True
|
||||
|
||||
async def declare_exchange(self, name, typ, durable=True):
|
||||
self.exchanges[name] = typ
|
||||
|
||||
async def get_exchange(self, name):
|
||||
return self.exchanges.setdefault(name, FakeExchange())
|
||||
|
||||
|
||||
class FakeConnection:
|
||||
def __init__(self):
|
||||
self.is_closed = False
|
||||
self._channel = FakeChannel()
|
||||
|
||||
async def channel(self):
|
||||
return self._channel
|
||||
|
||||
async def close(self):
|
||||
self.is_closed = True
|
||||
|
||||
|
||||
async def fake_connect_robust(url, **_):
|
||||
return FakeConnection()
|
||||
|
||||
|
||||
# ---------- tests ----------
|
||||
|
||||
|
||||
def test_publish_success(monkeypatch):
|
||||
_reset()
|
||||
monkeypatch.setattr(amqp, "HAS_AIO_PIKA", True)
|
||||
monkeypatch.setattr("aio_pika.connect_robust", fake_connect_robust)
|
||||
|
||||
asyncio.run(amqp.connect())
|
||||
ch = asyncio.run(amqp.get_channel())
|
||||
ex = FakeExchange()
|
||||
ch.exchanges[AMQP_EXCHANGE_ADMIN] = ex
|
||||
|
||||
asyncio.run(amqp.publish(AMQP_EXCHANGE_ADMIN, "test.key", {"foo": "bar"}))
|
||||
|
||||
assert len(ex.messages) == 1
|
||||
msg, rk = ex.messages[0]
|
||||
assert rk == "test.key"
|
||||
assert msg.body == b'{"foo": "bar"}'
|
||||
|
||||
|
||||
def test_publish_disconnected_no_raise(caplog):
|
||||
_reset()
|
||||
caplog.set_level(logging.ERROR)
|
||||
asyncio.run(amqp.publish(AMQP_EXCHANGE_ADMIN, "test.key", {"x": 1}))
|
||||
|
||||
assert len(caplog.records) > 0
|
||||
assert any("cannot publish" in r.message for r in caplog.records)
|
||||
|
||||
|
||||
def test_get_channel_reconnects_when_none(monkeypatch):
|
||||
_reset()
|
||||
monkeypatch.setattr(amqp, "HAS_AIO_PIKA", True)
|
||||
monkeypatch.setattr("aio_pika.connect_robust", fake_connect_robust)
|
||||
|
||||
ch = asyncio.run(amqp.get_channel())
|
||||
assert ch is not None
|
||||
assert not ch.is_closed
|
||||
Reference in New Issue
Block a user