diff --git a/README.md b/README.md index 589d137..c0c1ed9 100644 --- a/README.md +++ b/README.md @@ -113,8 +113,9 @@ If authentication is enabled (see below), send the taskowl API key as | **Tasks** | `list_tasks`, `get_task`, `get_task_timeline`, `get_task_chain`, `get_task_summary`, `list_task_types`, `list_orphaned_tasks` | | **Task actions** | `revoke_task`, `retry_task` | | **Workers** | `get_worker_status`, `list_workers`, `get_worker_stats`, `shutdown_worker`, `scale_worker_pool`, `get_active_tasks` | +| **Queues** | `list_queues` | -**Total: 15 tools** +**Total: 16 tools** `list_tasks` supports exact filters (`state`, `name`, `worker`, `since`), a partial case-insensitive `search` on the task name, `offset` for pagination, and `sort_by` @@ -133,6 +134,7 @@ Questions you can ask your AI assistant when the MCP server is connected: | "What's the retry chain for task abc?" | `get_task_chain` | | "What's the task success rate in the last 30 minutes?" | `get_task_summary` | | "Which workers are online?" | `get_worker_status`, `list_workers` | +| "How many messages are in each queue?" | `list_queues` | | "Shutdown worker celery@worker1" | `shutdown_worker` | | "Retry task abc" | `retry_task` | @@ -261,6 +263,7 @@ retries, and metrics. Interactive docs are available at: | **Task actions** | `POST /api/tasks/{id}/revoke`, `POST /api/tasks/{id}/retry` | | **Workers** | `GET /api/workers`, `GET /api/workers/list`, `GET /api/workers/{name}/stats`, `GET /api/workers/active-tasks` | | **Worker actions** | `POST /api/workers/{name}/shutdown`, `POST /api/workers/{name}/scale` | +| **Queues** | `GET /api/queues` | | **Ops** | `GET /health`, `GET /metrics` | The `/openapi.json` schema is the authoritative reference — this README lists diff --git a/src/taskowl/main.py b/src/taskowl/main.py index 11fee93..7e15298 100644 --- a/src/taskowl/main.py +++ b/src/taskowl/main.py @@ -24,6 +24,7 @@ list_task_types_query, list_tasks_query, ) +from taskowl.queues import list_queues from taskowl.workers import ( get_active_tasks, get_worker_stats, @@ -314,6 +315,17 @@ async def api_get_active_tasks( return await get_active_tasks(worker_name) +@app.get("/api/queues") +async def api_list_queues( + _: None = Depends(verify_api_key), +) -> dict: + """List Celery broker queues with message and consumer counts.""" + result = await list_queues() + if "error" in result: + raise HTTPException(status_code=502, detail=result["error"]) + return result + + # MCP server is now run separately via taskowl-mcp command # See taskowl.mcp.server:run_mcp_server() diff --git a/src/taskowl/mcp/tools.py b/src/taskowl/mcp/tools.py index e62625b..d9b86df 100644 --- a/src/taskowl/mcp/tools.py +++ b/src/taskowl/mcp/tools.py @@ -327,3 +327,17 @@ async def get_active_tasks(worker_name: str | None = None) -> dict: ) response.raise_for_status() return response.json() + + @server.tool( + name="list_queues", + description="List Celery broker queues with message and consumer counts", + ) + async def list_queues() -> dict: + """List Celery broker queues with message and consumer counts.""" + async with httpx.AsyncClient() as client: + response = await client.get( + f"http://{settings.taskowl_host}:{settings.taskowl_port}/api/queues", + headers=_get_headers(), + ) + response.raise_for_status() + return response.json() diff --git a/src/taskowl/queues.py b/src/taskowl/queues.py new file mode 100644 index 0000000..86f9df5 --- /dev/null +++ b/src/taskowl/queues.py @@ -0,0 +1,77 @@ +"""Broker queue monitoring functions for taskowl. + +This module provides functions to inspect Celery broker queues, +including message counts and consumer counts. It works with any +kombu-supported broker (RabbitMQ, Redis, ...). +""" + +import asyncio +import logging + +from celery import Celery +from kombu import Connection, Queue + +from taskowl.config import settings + +logger = logging.getLogger(__name__) + + +def _get_celery_app() -> Celery: + """Get Celery app instance configured with broker URL.""" + return Celery(broker=settings.celery_broker_url) + + +def _declared_queues(app: Celery) -> list[str]: + """Return the queue names declared by the Celery app. + + Falls back to the app's default queue name if the app exposes no + explicit queue configuration. + """ + queues = getattr(app.amqp, "queues", None) + if queues: + return list(queues.keys()) + return [app.conf.task_default_queue or "celery"] + + +def _list_queues_sync() -> dict: + """Synchronously inspect broker queues and return queue statistics.""" + app = _get_celery_app() + queue_names = _declared_queues(app) + + with Connection(settings.celery_broker_url) as conn: + channel = conn.channel() + queue_stats: list[dict] = [] + for name in queue_names: + declared = Queue(name).queue_declare(channel=channel, passive=True) + messages: int = int(getattr(declared, "message_count", 0)) + consumers: int = int(getattr(declared, "consumer_count", 0)) + queue_stats.append( + { + "name": name, + "messages": messages, + "consumers": consumers, + } + ) + + queue_stats.sort(key=lambda q: q["messages"], reverse=True) + total_messages: int = sum(int(q["messages"]) for q in queue_stats) + return { + "queues": queue_stats, + "total_messages": total_messages, + } + + +async def list_queues() -> dict: + """List Celery broker queues with message and consumer counts. + + Returns: + Dict with a queue list and a total message count. On broker + failure, returns an error dict. + """ + try: + # Broker I/O is synchronous; offload to a thread to keep the + # event loop responsive. + return await asyncio.to_thread(_list_queues_sync) + except Exception as e: + logger.exception("Failed to list broker queues") + return {"error": f"Failed to list queues: {str(e)}"} diff --git a/tests/test_api.py b/tests/test_api.py index c7b6c4c..bea2a46 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -848,6 +848,31 @@ async def test_api_get_active_tasks_with_worker_filter(client: AsyncClient): assert "active_tasks" in data +@pytest.mark.asyncio +async def test_api_list_queues(client: AsyncClient): + """Test GET /api/queues.""" + with patch("taskowl.queues._list_queues_sync") as mock_list_queues: + mock_list_queues.return_value = { + "queues": [{"name": "celery", "messages": 3, "consumers": 1}], + "total_messages": 3, + } + + response = await client.get("/api/queues") + assert response.status_code == 200 + data = response.json() + assert data["total_messages"] == 3 + assert data["queues"][0]["name"] == "celery" + + +@pytest.mark.asyncio +async def test_api_list_queues_error(client: AsyncClient): + """Test GET /api/queues when the broker fails.""" + with patch("taskowl.queues._list_queues_sync", side_effect=Exception("Broker down")): + response = await client.get("/api/queues") + assert response.status_code == 502 + assert "Broker down" in response.json()["detail"] + + @pytest.mark.asyncio async def test_api_list_orphaned_tasks_empty(client: AsyncClient): """Test GET /api/tasks/orphaned with no orphaned tasks.""" diff --git a/tests/test_queues.py b/tests/test_queues.py new file mode 100644 index 0000000..fb43a48 --- /dev/null +++ b/tests/test_queues.py @@ -0,0 +1,129 @@ +"""Tests for broker queue monitoring functions.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from taskowl.queues import list_queues + + +def _make_declared(message_count=0, consumer_count=0): + declared = MagicMock() + declared.message_count = message_count + declared.consumer_count = consumer_count + return declared + + +@pytest.mark.asyncio +async def test_list_queues_success(): + """Test listing queues with message and consumer counts.""" + mock_connection = MagicMock() + mock_channel = MagicMock() + + # Two queues: default with messages, secondary empty + mock_connection.channel.return_value = mock_channel + + declared_default = _make_declared(message_count=5, consumer_count=1) + declared_secondary = _make_declared(message_count=0, consumer_count=0) + + with ( + patch("taskowl.queues.Connection", return_value=mock_connection), + patch("taskowl.queues._get_celery_app") as mock_get_app, + patch("taskowl.queues.Queue") as mock_queue_cls, + ): + mock_app = MagicMock() + mock_app.amqp.queues.keys.return_value = ["default", "secondary"] + mock_get_app.return_value = mock_app + + def make_queue(name): + q = MagicMock() + q.queue_declare.return_value = ( + declared_default if name == "default" else declared_secondary + ) + return q + + mock_queue_cls.side_effect = make_queue + + result = await list_queues() + + assert "queues" in result + queues = result["queues"] + assert len(queues) == 2 + assert queues[0]["name"] == "default" + assert queues[0]["messages"] == 5 + assert queues[0]["consumers"] == 1 + assert queues[1]["name"] == "secondary" + assert queues[1]["messages"] == 0 + assert result["total_messages"] == 5 + + +@pytest.mark.asyncio +async def test_list_queues_ordered_by_messages_desc(): + """Test queues are ordered by message count descending.""" + mock_connection = MagicMock() + mock_connection.channel.return_value = MagicMock() + + with ( + patch("taskowl.queues.Connection", return_value=mock_connection), + patch("taskowl.queues._get_celery_app") as mock_get_app, + patch("taskowl.queues.Queue") as mock_queue_cls, + ): + mock_app = MagicMock() + mock_app.amqp.queues.keys.return_value = ["small", "big", "mid"] + mock_get_app.return_value = mock_app + + counts = {"small": 1, "big": 10, "mid": 5} + + def make_queue(name): + q = MagicMock() + q.queue_declare.return_value = _make_declared( + message_count=counts[name], consumer_count=0 + ) + return q + + mock_queue_cls.side_effect = make_queue + + result = await list_queues() + + names = [q["name"] for q in result["queues"]] + assert names == ["big", "mid", "small"] + + +@pytest.mark.asyncio +async def test_list_queues_error(): + """Test listing queues when the broker fails.""" + with patch("taskowl.queues.Connection", side_effect=Exception("Broker down")): + result = await list_queues() + + assert "error" in result + assert "Broker down" in result["error"] + + +@pytest.mark.asyncio +async def test_list_queues_no_queues_configured(): + """Test listing queues when the app exposes no explicit queue config.""" + mock_connection = MagicMock() + mock_connection.channel.return_value = MagicMock() + + with ( + patch("taskowl.queues.Connection", return_value=mock_connection), + patch("taskowl.queues._get_celery_app") as mock_get_app, + patch("taskowl.queues.Queue") as mock_queue_cls, + ): + mock_app = MagicMock() + # Simulate an app without explicit queue config -> use default queue + mock_app.amqp.queues = None + mock_app.conf.task_default_queue = "celery" + mock_get_app.return_value = mock_app + + declared = _make_declared(message_count=3, consumer_count=1) + mock_queue = MagicMock() + mock_queue.queue_declare.return_value = declared + mock_queue_cls.return_value = mock_queue + + result = await list_queues() + + assert len(result["queues"]) == 1 + assert result["queues"][0]["name"] == "celery" + assert result["queues"][0]["messages"] == 3 + assert result["total_messages"] == 3