OpenJarvis/tests/server/test_agent_manager_routes.py

"""Tests for Agent Manager API routes."""

from __future__ import annotations

import json
import tempfile
from pathlib import Path
from unittest.mock import MagicMock

import pytest

from openjarvis.agents.manager import AgentManager


@pytest.fixture
def manager():
    with tempfile.TemporaryDirectory() as tmpdir:
        mgr = AgentManager(db_path=str(Path(tmpdir) / "agents.db"))
        yield mgr
        mgr.close()


try:
    from fastapi.testclient import TestClient

    HAS_FASTAPI = True
except ImportError:
    HAS_FASTAPI = False


@pytest.mark.skipif(not HAS_FASTAPI, reason="fastapi not installed")
class TestAgentManagerRoutes:
    @pytest.fixture
    def client(self, manager):
        from fastapi import FastAPI

        from openjarvis.server.agent_manager_routes import create_agent_manager_router

        app = FastAPI()
        routers = create_agent_manager_router(manager)
        agents_router, templates_router, global_router, tools_router = routers
        app.include_router(agents_router)
        app.include_router(templates_router)
        app.include_router(global_router)
        app.include_router(tools_router)
        return TestClient(app)

    def test_list_agents_empty(self, client):
        resp = client.get("/v1/managed-agents")
        assert resp.status_code == 200
        assert resp.json()["agents"] == []

    def test_create_agent(self, client):
        resp = client.post(
            "/v1/managed-agents",
            json={
                "name": "researcher",
                "agent_type": "monitor_operative",
            },
        )
        assert resp.status_code == 200
        data = resp.json()
        assert data["name"] == "researcher"
        assert data["status"] == "idle"

    def test_get_agent(self, client):
        create_resp = client.post("/v1/managed-agents", json={"name": "test"})
        agent_id = create_resp.json()["id"]
        resp = client.get(f"/v1/managed-agents/{agent_id}")
        assert resp.status_code == 200
        assert resp.json()["id"] == agent_id

    def test_get_agent_not_found(self, client):
        resp = client.get("/v1/managed-agents/nonexistent")
        assert resp.status_code == 404

    def test_update_agent(self, client):
        create_resp = client.post("/v1/managed-agents", json={"name": "old"})
        agent_id = create_resp.json()["id"]
        resp = client.patch(f"/v1/managed-agents/{agent_id}", json={"name": "new"})
        assert resp.status_code == 200
        assert resp.json()["name"] == "new"

    def test_delete_agent(self, client):
        create_resp = client.post("/v1/managed-agents", json={"name": "doomed"})
        agent_id = create_resp.json()["id"]
        resp = client.delete(f"/v1/managed-agents/{agent_id}")
        assert resp.status_code == 200

    def test_pause_resume(self, client):
        create_resp = client.post("/v1/managed-agents", json={"name": "pausable"})
        agent_id = create_resp.json()["id"]
        client.post(f"/v1/managed-agents/{agent_id}/pause")
        resp = client.get(f"/v1/managed-agents/{agent_id}")
        assert resp.json()["status"] == "paused"
        client.post(f"/v1/managed-agents/{agent_id}/resume")
        resp = client.get(f"/v1/managed-agents/{agent_id}")
        assert resp.json()["status"] == "idle"

    def test_create_task(self, client):
        create_resp = client.post("/v1/managed-agents", json={"name": "worker"})
        agent_id = create_resp.json()["id"]
        resp = client.post(
            f"/v1/managed-agents/{agent_id}/tasks",
            json={
                "description": "Find papers on reasoning",
            },
        )
        assert resp.status_code == 200
        assert resp.json()["description"] == "Find papers on reasoning"

    def test_list_tasks(self, client):
        create_resp = client.post("/v1/managed-agents", json={"name": "worker"})
        agent_id = create_resp.json()["id"]
        client.post(f"/v1/managed-agents/{agent_id}/tasks", json={"description": "t1"})
        client.post(f"/v1/managed-agents/{agent_id}/tasks", json={"description": "t2"})
        resp = client.get(f"/v1/managed-agents/{agent_id}/tasks")
        assert len(resp.json()["tasks"]) == 2

    def test_channel_binding_crud(self, client):
        create_resp = client.post("/v1/managed-agents", json={"name": "slacker"})
        agent_id = create_resp.json()["id"]
        # Bind
        bind_resp = client.post(
            f"/v1/managed-agents/{agent_id}/channels",
            json={
                "channel_type": "slack",
                "config": {"channel": "#research"},
            },
        )
        assert bind_resp.status_code == 200
        binding_id = bind_resp.json()["id"]
        # List
        list_resp = client.get(f"/v1/managed-agents/{agent_id}/channels")
        assert len(list_resp.json()["bindings"]) == 1
        # Unbind
        url = f"/v1/managed-agents/{agent_id}/channels/{binding_id}"
        unbind_resp = client.delete(url)
        assert unbind_resp.status_code == 200

    def test_templates(self, client):
        resp = client.get("/v1/templates")
        assert resp.status_code == 200
        templates = resp.json()["templates"]
        assert any(t["id"] == "research_monitor" for t in templates)

    def test_recover_agent(self, manager, client):
        # Create agent, save checkpoint, set error status
        agent = manager.create_agent(name="err", agent_type="simple")
        manager.save_checkpoint(agent["id"], "tick-1", {"msgs": []}, {})
        manager.update_agent(agent["id"], status="error")

        res = client.post(f"/v1/managed-agents/{agent['id']}/recover")
        assert res.status_code == 200
        body = res.json()
        assert body["recovered"] is True
        assert body["checkpoint"]["tick_id"] == "tick-1"

    def test_recover_agent_no_checkpoint(self, manager, client):
        agent = manager.create_agent(name="err", agent_type="simple")
        manager.update_agent(agent["id"], status="error")
        res = client.post(f"/v1/managed-agents/{agent['id']}/recover")
        assert res.status_code == 200
        body = res.json()
        assert body["recovered"] is True
        assert body["checkpoint"] is None
        # Status should be reset to idle
        refreshed = manager.get_agent(agent["id"])
        assert refreshed["status"] == "idle"

    def test_list_error_agents(self, manager, client):
        manager.create_agent(name="ok", agent_type="simple")
        err = manager.create_agent(name="broken", agent_type="simple")
        manager.update_agent(err["id"], status="error")

        res = client.get("/v1/agents/errors")
        assert res.status_code == 200
        agents = res.json()["agents"]
        assert len(agents) == 1
        assert agents[0]["name"] == "broken"

    def test_send_and_list_messages(self, manager, client):
        agent = manager.create_agent(name="chat", agent_type="simple")

        res = client.post(
            f"/v1/managed-agents/{agent['id']}/messages",
            json={"content": "hello", "mode": "queued"},
        )
        assert res.status_code == 200

        res = client.get(f"/v1/managed-agents/{agent['id']}/messages")
        assert res.status_code == 200
        assert len(res.json()["messages"]) == 1

    def test_get_agent_state(self, manager, client):
        agent = manager.create_agent(name="stateful", agent_type="simple")
        res = client.get(f"/v1/managed-agents/{agent['id']}/state")
        assert res.status_code == 200
        state = res.json()
        assert "agent" in state
        assert "tasks" in state
        assert "channels" in state
        assert "messages" in state
        assert "checkpoint" in state

    def test_send_message_non_stream_unchanged(self, manager, client):
        """stream=False (default) returns a normal JSON message, not SSE."""
        agent = manager.create_agent(name="basic", agent_type="simple")
        res = client.post(
            f"/v1/managed-agents/{agent['id']}/messages",
            json={"content": "hello", "stream": False},
        )
        assert res.status_code == 200
        data = res.json()
        assert data["content"] == "hello"
        assert data["direction"] == "user_to_agent"

    def test_send_message_stream_not_found(self, manager, client):
        """Streaming to a non-existent agent returns 404."""
        res = client.post(
            "/v1/managed-agents/nonexistent/messages",
            json={"content": "hello", "stream": True},
        )
        assert res.status_code == 404


def test_run_agent_concurrent_returns_409(tmp_path):
    """Rapid Run Now clicks should not spawn multiple ticks."""
    from openjarvis.agents.manager import AgentManager

    mgr = AgentManager(db_path=str(tmp_path / "test.db"))
    agent = mgr.create_agent("Test", config={"schedule_type": "manual"})
    aid = agent["id"]

    # Simulate first click acquiring the tick
    mgr.start_tick(aid)

    # Second click should fail
    with pytest.raises(ValueError, match="already executing a tick"):
        mgr.start_tick(aid)

    mgr.end_tick(aid)


@pytest.mark.skipif(not HAS_FASTAPI, reason="fastapi not installed")
class TestAgentManagerStreaming:
    """Tests for the SSE streaming mode of the managed-agent messages endpoint."""

    @pytest.fixture
    def _mock_engine(self):
        engine = MagicMock()
        engine.engine_id = "mock"
        engine._model = "test-model"
        engine.health.return_value = True
        return engine

    @pytest.fixture
    def _mock_agent_cls(self):
        """Register a mock agent class in the AgentRegistry for testing."""
        from openjarvis.agents._stubs import AgentResult
        from openjarvis.core.registry import AgentRegistry

        class _MockStreamAgent:
            agent_id = "mock_stream"

            def __init__(self, engine, model, **kwargs):
                self._engine = engine
                self._model = model

            def run(self, input_text, context=None, **kwargs):
                return AgentResult(content=f"Echo: {input_text}", turns=1)

        # Register under a unique key for test isolation
        AgentRegistry._entries()["_test_stream"] = _MockStreamAgent
        yield _MockStreamAgent
        AgentRegistry._entries().pop("_test_stream", None)

    @pytest.fixture
    def stream_client(self, manager, _mock_engine, _mock_agent_cls):
        from fastapi import FastAPI

        from openjarvis.server.agent_manager_routes import create_agent_manager_router

        app = FastAPI()
        app.state.engine = _mock_engine
        app.state.bus = None

        routers = create_agent_manager_router(manager)
        agents_router, templates_router, global_router, tools_router = routers
        app.include_router(agents_router)
        app.include_router(templates_router)
        app.include_router(global_router)
        app.include_router(tools_router)
        return TestClient(app)

    def test_send_message_stream(self, manager, stream_client, _mock_agent_cls):
        """Test streaming mode returns SSE response with [DONE] sentinel."""
        agent = manager.create_agent(
            name="streamer", agent_type="_test_stream",
        )
        resp = stream_client.post(
            f"/v1/managed-agents/{agent['id']}/messages",
            json={"content": "What is 2+2?", "stream": True},
        )
        assert resp.status_code == 200
        assert "text/event-stream" in resp.headers.get("content-type", "")

        # Parse SSE events
        lines = resp.text.strip().split("\n")
        data_lines = [ln for ln in lines if ln.startswith("data:")]
        assert len(data_lines) > 0
        # Last data line must be [DONE]
        assert data_lines[-1].strip() == "data: [DONE]"

    def test_send_message_stream_content(self, manager, stream_client, _mock_agent_cls):
        """Test streaming returns the correct agent response content."""
        agent = manager.create_agent(
            name="streamer2", agent_type="_test_stream",
        )
        resp = stream_client.post(
            f"/v1/managed-agents/{agent['id']}/messages",
            json={"content": "Hello world", "stream": True},
        )
        assert resp.status_code == 200

        # Collect content tokens from stream
        content = ""
        for line in resp.text.strip().split("\n"):
            if line.startswith("data:") and "[DONE]" not in line:
                raw = line[5:].strip()
                try:
                    data = json.loads(raw)
                except json.JSONDecodeError:
                    continue
                choices = data.get("choices", [{}])
                delta_content = choices[0].get("delta", {}).get("content")
                if delta_content:
                    content += delta_content

        assert content == "Echo: Hello world"

    def test_send_message_stream_stores_response(
        self, manager, stream_client, _mock_agent_cls,
    ):
        """After streaming, agent response is persisted in the DB."""
        agent = manager.create_agent(
            name="streamer3", agent_type="_test_stream",
        )
        resp = stream_client.post(
            f"/v1/managed-agents/{agent['id']}/messages",
            json={"content": "persist me", "stream": True},
        )
        assert resp.status_code == 200

        # Check messages in DB
        messages = manager.list_messages(agent["id"])
        # Should have both the user message and the agent response
        assert len(messages) == 2
        directions = {m["direction"] for m in messages}
        assert "user_to_agent" in directions
        assert "agent_to_user" in directions
        agent_msg = next(m for m in messages if m["direction"] == "agent_to_user")
        assert "persist me" in agent_msg["content"]

    def test_send_message_stream_finish_reason(
        self, manager, stream_client, _mock_agent_cls,
    ):
        """The final chunk before [DONE] has finish_reason='stop'."""
        agent = manager.create_agent(
            name="streamer4", agent_type="_test_stream",
        )
        resp = stream_client.post(
            f"/v1/managed-agents/{agent['id']}/messages",
            json={"content": "check finish", "stream": True},
        )
        # Collect all data chunks (excluding [DONE])
        chunks = []
        for line in resp.text.strip().split("\n"):
            if line.startswith("data:") and "[DONE]" not in line:
                raw = line[5:].strip()
                try:
                    chunks.append(json.loads(raw))
                except json.JSONDecodeError:
                    continue

        # Last chunk should have finish_reason="stop"
        assert chunks[-1]["choices"][0]["finish_reason"] == "stop"