feat: Matrix E2EE key management + multi-user isolation

- Add rag_key_manager.py: stores encryption key in private E2EE room
- Bot loads key from Matrix on startup, injects into RAG via portal proxy
- No plaintext key on disk (removed RAG_ENCRYPTION_KEY from .env)
- Pass owner_id (matrix_user_id) to RAG search for user isolation
- Stronger format_context instructions for source link rendering

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Christian Gick
2026-03-03 11:19:02 +00:00
parent 5d3a6c8c79
commit 9578e0406b
2 changed files with 215 additions and 5 deletions

34
bot.py
View File

@@ -15,6 +15,7 @@ import fitz # pymupdf
import httpx
from openai import AsyncOpenAI
from olm import sas as olm_sas
from rag_key_manager import RAGKeyManager
from nio import (
AsyncClient,
@@ -335,15 +336,17 @@ class DocumentRAG:
# Prefer customer-VM RAG service (encrypted, local)
if self.use_local_rag:
return await self._search_local(query, top_k)
return await self._search_local(query, top_k, matrix_user_id)
# Fallback: central portal API (legacy, unencrypted)
return await self._search_portal(query, top_k, matrix_user_id)
async def _search_local(self, query: str, top_k: int) -> list[dict]:
async def _search_local(self, query: str, top_k: int, matrix_user_id: str | None = None) -> list[dict]:
"""Search via customer-VM RAG service (localhost)."""
try:
body = {"query": query, "limit": top_k}
if matrix_user_id:
body["owner_id"] = matrix_user_id
headers: dict[str, str] = {"Content-Type": "application/json"}
if self.rag_auth_token:
headers["Authorization"] = f"Bearer {self.rag_auth_token}"
@@ -415,9 +418,13 @@ class DocumentRAG:
parts.append(f"Content:\n{content}")
parts.append("") # blank line between docs
parts.append("Use the document content above to answer the user's question. "
"When referencing documents, use markdown links: [Document Title](url). "
"Never show raw URLs.")
parts.append("IMPORTANT INSTRUCTIONS FOR DOCUMENT RESPONSES:\n"
"1. Answer the user's question using the document content above.\n"
"2. You MUST include a source link for EVERY document you reference.\n"
"3. Format links as markdown: [Document Title](url)\n"
"4. Place the link right after mentioning or quoting the document.\n"
"5. If a document has no link, skip the link but still reference the title.\n"
"6. Never show raw URLs without markdown formatting.")
return "\n".join(parts)
@@ -940,6 +947,7 @@ class Bot:
self.active_callers: dict[str, set[str]] = {} # room_id → set of caller user IDs
self.rag = DocumentRAG(PORTAL_URL, BOT_API_KEY,
rag_endpoint=RAG_ENDPOINT, rag_auth_token=RAG_AUTH_TOKEN)
self.key_manager = RAGKeyManager(self.client, PORTAL_URL, BOT_API_KEY)
self.memory = MemoryClient(MEMORY_SERVICE_URL)
self.atlassian = AtlassianClient(PORTAL_URL, BOT_API_KEY)
self.llm = AsyncOpenAI(base_url=LITELLM_URL, api_key=LITELLM_KEY) if LITELLM_URL else None
@@ -1035,6 +1043,20 @@ class Bot:
await self.client.sync_forever(timeout=30000, full_state=True)
async def _inject_rag_key(self):
"""Load document encryption key from Matrix and inject into RAG service."""
try:
seed_key = os.environ.get("RAG_ENCRYPTION_KEY_SEED")
success = await self.key_manager.ensure_rag_key(seed_key_hex=seed_key)
if success:
logger.info("RAG encryption key loaded from Matrix E2EE")
if seed_key:
logger.info("Migration complete - RAG_ENCRYPTION_KEY_SEED can now be removed from env")
else:
logger.warning("Failed to load RAG encryption key - documents will be inaccessible")
except Exception as e:
logger.error("RAG key injection failed: %s", e, exc_info=True)
async def on_invite(self, room, event: InviteMemberEvent):
if event.state_key != BOT_USER:
return
@@ -1046,6 +1068,8 @@ class Bot:
if not self._sync_token_received:
self._sync_token_received = True
logger.info("Initial sync complete, text handler active")
# Inject RAG encryption key from Matrix E2EE room
asyncio.create_task(self._inject_rag_key())
for user_id in list(self.client.device_store.users):
for device in self.client.device_store.active_user_devices(user_id):
if not device.verified: