import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from src.services.api.middlewares.auth import BrainPATMiddleware from src.services.api.middlewares.brains import BrainMiddleware from src.services.api.routes.agent import agent_router from tests.test_workspace_facade import FacadeCache, FacadeData def agent_client(monkeypatch): from src.services.api.middlewares import auth, brains from src.services.api.routes import agent data = FacadeData() cache = FacadeCache() monkeypatch.setenv("BRAINPAT_TOKEN", "system-secret") monkeypatch.setenv("PUBLIC_BASE_URL", "https://brain.example") monkeypatch.setattr(brains, "data_adapter", data) monkeypatch.setattr(auth, "data_adapter", data) monkeypatch.setattr(agent, "data_adapter", data) calls = [] async def search(query, limit, brain_id): calls.append(("search", brain_id, query)) return [{"id": brain_id, "type": "chunk", "text": query}] async def context(query, max_tokens, brain_id): calls.append(("context", brain_id, query)) return {"context": f"{brain_id}:{query}", "sources": []} async def store(**kwargs): calls.append( ( "store", kwargs["brain_id"], kwargs["text"], kwargs["metadata"], kwargs["request"].headers.get("Idempotency-Key"), ) ) return {"task_id": f'{kwargs["brain_id"]}-task'} async def task(task_id, brain_id): if task_id != f"{brain_id}-task": from fastapi import HTTPException raise HTTPException(status_code=404, detail="Task not found") return { "task_id": task_id, "brain_id": brain_id, "status": "queued", "workspace": "untrusted-task-payload", } monkeypatch.setattr(agent, "memory_search", search) monkeypatch.setattr(agent, "memory_context", context) monkeypatch.setattr(agent, "memory_store", store) monkeypatch.setattr(agent, "memory_task_status", task) app = FastAPI() app.add_middleware(BrainPATMiddleware) app.add_middleware(BrainMiddleware) app.include_router(agent_router) return TestClient(app), calls def test_agent_operations_use_canonical_workspace_scope(monkeypatch): client, calls = agent_client(monkeypatch) headers = {"Authorization": "Bearer system-secret"} assert client.post( "/brains/alpha/agent/search", headers=headers, json={"query": "memory"} ).json()["results"][0]["id"] == "alpha" assert client.post( "/brains/beta/agent/context", headers=headers, json={"query": "context"} ).json()["context"] == "beta:context" stored = client.post( "/brains/alpha/agent/memory", headers={**headers, "Idempotency-Key": "agent-message-1"}, json={"text": "ALPHA_ONLY_FACT", "metadata": {"source": "hermes"}}, ) assert stored.status_code == 202 assert stored.json() == { "workspace": "alpha", "status": "queued", "task_id": "alpha-task", } assert ( "store", "alpha", "ALPHA_ONLY_FACT", {"source": "hermes"}, "agent-message-1", ) in calls assert client.get( "/brains/beta/agent/tasks/alpha-task", headers=headers ).status_code == 404 own_task = client.get("/brains/alpha/agent/tasks/alpha-task", headers=headers) assert own_task.json()["workspace"] == "alpha" def test_agent_security_discovery_and_openapi(monkeypatch): client, _ = agent_client(monkeypatch) system = {"Authorization": "Bearer system-secret"} assert client.get( "/brains/beta/agent/capabilities", headers={"Authorization": "Bearer alpha-secret"}, ).status_code == 403 assert client.get( "/brains/beta/agent/capabilities", headers={"Authorization": "Bearer invalid-secret"}, ).status_code == 401 assert client.get("/brains/archived/agent/capabilities", headers=system).status_code == 410 assert client.get("/brains/missing/agent/capabilities", headers=system).status_code == 404 conflict = client.post( "/brains/alpha/agent/search", headers=system, json={"query": "x", "brain_id": "beta"}, ) assert conflict.status_code == 409 capabilities = client.get( "/brains/alpha/agent/capabilities", headers=system ).json() assert capabilities["workspace"] == { "slug": "alpha", "brain_id": "alpha", "display_name": "Alpha", } assert capabilities["endpoints"]["agent_api"].endswith("/brains/alpha/agent") assert capabilities["endpoints"]["native_openapi"].endswith( "/brains/alpha/api/openapi.json" ) assert capabilities["endpoints"]["agent_openapi"].endswith( "/brains/alpha/agent/openapi.json" ) assert capabilities["endpoints"]["mcp"].endswith("/brains/alpha/mcp") schema = client.get("/brains/alpha/agent/openapi.json", headers=system).json() assert schema["servers"] == [{"url": "https://brain.example/brains/alpha/agent"}] assert set(schema["paths"]) == { "/search", "/context", "/memory", "/tasks/{task_id}", "/capabilities", } def test_agent_metadata_survives_canonical_ingestion_and_chunk_retrieval(monkeypatch): import asyncio from unittest.mock import MagicMock, patch from fastapi.responses import JSONResponse from src.constants.embeddings import Vector from src.services import agent_memory from src.services.api.controllers import structured_data from src.workers.tasks import ingestion queued_payload = {} async def queue(body, request, brain_id): body.brain_id = brain_id queued_payload.update(body.model_dump()) return JSONResponse({"message": "accepted", "task_id": "alpha-task"}, 202) monkeypatch.setattr(agent_memory, "ingest_data", queue) result = asyncio.run( agent_memory.memory_store( text="Customer prefers annual billing.", metadata={"source": "hermes", "session_id": "abc"}, brain_id="alpha", request=type("Request", (), {"headers": {}})(), ) ) assert result["task_id"] == "alpha-task" assert queued_payload["brain_id"] == "alpha" assert queued_payload["meta_keys"] == { "source": "hermes", "session_id": "abc", } class MemoryData: def __init__(self): self.chunks = {} def save_text_chunk(self, chunk, brain_id): self.chunks.setdefault(brain_id, []).append(chunk) return chunk def get_text_chunks( self, brain_id, limit, skip, query_text, metadata_eq, order ): chunks = self.chunks.get(brain_id, []) if metadata_eq: chunks = [ chunk for chunk in chunks if all((chunk.metadata or {}).get(k) == v for k, v in metadata_eq.items()) ] return chunks[skip : skip + limit], len(chunks) repository = MemoryData() request = MagicMock(id="alpha-task") with patch.object(ingestion, "data_adapter", repository), patch.object( ingestion, "embeddings_adapter" ) as embeddings, patch.object( ingestion, "vector_store_adapter" ), patch.object( ingestion.config, "pipeline_mode", "lightweight" ), patch.object( ingestion, "set_ingestion_task_status" ), patch.object( structured_data, "data_adapter", repository ): embeddings.embed_text.return_value = Vector( id="vector-1", embeddings=[0.1], metadata={} ) queued_payload["skip_enrichment"] = True ingestion.ingest_data.run.__func__( type("Bound", (), {"request": request})(), queued_payload ) chunks, total = asyncio.run( structured_data.get_text_chunks( "alpha", metadata_eq={"source": "hermes", "session_id": "abc"} ) ) assert total == 1 assert chunks[0].metadata == {"source": "hermes", "session_id": "abc"} def test_agent_memory_idempotency_key_reuses_task_identifier(monkeypatch): import asyncio from src.services import agent_memory from src.services.api.routes import ingest queued = [] monkeypatch.setattr( ingest.ingest_data_task, "apply_async", lambda *, args, task_id: queued.append((args[0]["brain_id"], task_id)), ) monkeypatch.setattr(ingest, "set_ingestion_task_status", lambda *args, **kwargs: {}) request = type("Request", (), {"headers": {"Idempotency-Key": "message-42"}})() first = asyncio.run( agent_memory.memory_store( text="Customer prefers annual billing.", metadata={"source": "hermes"}, brain_id="alpha", request=request, ) ) second = asyncio.run( agent_memory.memory_store( text="Customer prefers annual billing.", metadata={"source": "hermes"}, brain_id="alpha", request=request, ) ) assert first["task_id"] == second["task_id"] == "message-42" assert queued == [("alpha", "message-42"), ("alpha", "message-42")] def test_relationship_worker_rejects_missing_brain_scope(): from src.workers.tasks import ingestion with pytest.raises(ValueError, match="requires brain_id"): ingestion.process_architect_relationships.run.__func__( type("Bound", (), {"request": type("Request", (), {"id": "task"})()})(), {"relationships": []}, )