tests/test_agent_api.py

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": []},
        )