diff --git a/README.md b/README.md index e97f103..e373738 100644 --- a/README.md +++ b/README.md @@ -112,10 +112,10 @@ 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`, `execute_task` | -| **Workers** | `get_worker_status`, `list_workers`, `get_worker_stats`, `shutdown_worker`, `scale_worker_pool`, `get_active_tasks` | +| **Workers** | `get_worker_status`, `list_workers`, `get_worker_stats`, `shutdown_worker`, `scale_worker_pool`, `get_active_tasks`, `get_scheduled_tasks`, `get_reserved_tasks` | | **Queues** | `list_queues` | -**Total: 17 tools** +**Total: 19 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` @@ -136,6 +136,7 @@ Questions you can ask your AI assistant when the MCP server is connected: | "Which workers are online?" | `get_worker_status`, `list_workers` | | "How many messages are in each queue?" | `list_queues` | | "Shutdown worker celery@worker1" | `shutdown_worker` | +| "What's scheduled to run next?" | `get_scheduled_tasks`, `get_reserved_tasks` | | "Retry task abc" | `retry_task` | | "Run myapp.tasks.process now" | `execute_task` | @@ -262,7 +263,7 @@ retries, and metrics. Interactive docs are available at: |------|-----------| | **Tasks** | `GET /api/tasks`, `GET /api/tasks/{id}`, `GET /api/tasks/{id}/timeline`, `GET /api/tasks/{id}/chain`, `GET /api/tasks/summary`, `GET /api/tasks/types`, `GET /api/tasks/orphaned` | | **Task actions** | `POST /api/tasks/{id}/revoke`, `POST /api/tasks/{id}/retry`, `POST /api/tasks/execute` | -| **Workers** | `GET /api/workers`, `GET /api/workers/list`, `GET /api/workers/{name}/stats`, `GET /api/workers/active-tasks` | +| **Workers** | `GET /api/workers`, `GET /api/workers/list`, `GET /api/workers/{name}/stats`, `GET /api/workers/active-tasks`, `GET /api/workers/scheduled`, `GET /api/workers/reserved` | | **Worker actions** | `POST /api/workers/{name}/shutdown`, `POST /api/workers/{name}/scale` | | **Queues** | `GET /api/queues` | | **Ops** | `GET /health`, `GET /metrics` | diff --git a/src/taskowl/main.py b/src/taskowl/main.py index de5c7a8..6e841e7 100644 --- a/src/taskowl/main.py +++ b/src/taskowl/main.py @@ -27,6 +27,8 @@ from taskowl.queues import list_queues from taskowl.workers import ( get_active_tasks, + get_reserved_tasks, + get_scheduled_tasks, get_worker_stats, list_workers, scale_worker_pool, @@ -353,6 +355,32 @@ async def api_get_active_tasks( return await get_active_tasks(worker_name) +@app.get("/api/workers/scheduled") +async def api_get_scheduled_tasks( + worker_name: str | None = None, + _: None = Depends(verify_api_key), +) -> dict: + """Get tasks scheduled to run (with an ETA/countdown). + + Args: + worker_name: Optional worker name to filter by + """ + return await get_scheduled_tasks(worker_name) + + +@app.get("/api/workers/reserved") +async def api_get_reserved_tasks( + worker_name: str | None = None, + _: None = Depends(verify_api_key), +) -> dict: + """Get tasks reserved (prefetched but not started). + + Args: + worker_name: Optional worker name to filter by + """ + return await get_reserved_tasks(worker_name) + + @app.get("/api/queues") async def api_list_queues( _: None = Depends(verify_api_key), diff --git a/src/taskowl/mcp/tools.py b/src/taskowl/mcp/tools.py index 4b459d2..fe9fcb9 100644 --- a/src/taskowl/mcp/tools.py +++ b/src/taskowl/mcp/tools.py @@ -379,6 +379,58 @@ async def get_active_tasks(worker_name: str | None = None) -> dict: response.raise_for_status() return response.json() + @server.tool( + name="get_scheduled_tasks", + description=( + "Get tasks scheduled to run (with an ETA/countdown) across all workers " + "or a specific worker" + ), + ) + async def get_scheduled_tasks(worker_name: str | None = None) -> dict: + """Get tasks scheduled to run. + + Args: + worker_name: Optional worker name to filter by + """ + async with httpx.AsyncClient() as client: + params = {} + if worker_name: + params["worker_name"] = worker_name + + response = await client.get( + f"http://{settings.taskowl_host}:{settings.taskowl_port}/api/workers/scheduled", + params=params, + headers=_get_headers(), + ) + response.raise_for_status() + return response.json() + + @server.tool( + name="get_reserved_tasks", + description=( + "Get tasks reserved (prefetched but not started) across all workers " + "or a specific worker" + ), + ) + async def get_reserved_tasks(worker_name: str | None = None) -> dict: + """Get tasks reserved (prefetched but not started). + + Args: + worker_name: Optional worker name to filter by + """ + async with httpx.AsyncClient() as client: + params = {} + if worker_name: + params["worker_name"] = worker_name + + response = await client.get( + f"http://{settings.taskowl_host}:{settings.taskowl_port}/api/workers/reserved", + params=params, + headers=_get_headers(), + ) + response.raise_for_status() + return response.json() + @server.tool( name="list_queues", description="List Celery broker queues with message and consumer counts", diff --git a/src/taskowl/workers.py b/src/taskowl/workers.py index b0f78c6..1c1a4c7 100644 --- a/src/taskowl/workers.py +++ b/src/taskowl/workers.py @@ -145,3 +145,49 @@ async def get_active_tasks(worker_name: str | None = None) -> dict: return {"active_tasks": active or {}} except Exception as e: return {"error": f"Failed to get active tasks: {str(e)}"} + + +async def get_scheduled_tasks(worker_name: str | None = None) -> dict: + """Get tasks scheduled to run (with an ETA/countdown) on workers. + + Args: + worker_name: Optional worker name to filter by + + Returns: + Dict with scheduled tasks grouped by worker + """ + try: + app = _get_celery_app() + inspect = app.control.inspect() + + if worker_name: + inspect = inspect.destination([worker_name]) + + scheduled = inspect.scheduled() + + return {"scheduled_tasks": scheduled or {}} + except Exception as e: + return {"error": f"Failed to get scheduled tasks: {str(e)}"} + + +async def get_reserved_tasks(worker_name: str | None = None) -> dict: + """Get tasks reserved (prefetched but not started) on workers. + + Args: + worker_name: Optional worker name to filter by + + Returns: + Dict with reserved tasks grouped by worker + """ + try: + app = _get_celery_app() + inspect = app.control.inspect() + + if worker_name: + inspect = inspect.destination([worker_name]) + + reserved = inspect.reserved() + + return {"reserved_tasks": reserved or {}} + except Exception as e: + return {"error": f"Failed to get reserved tasks: {str(e)}"} diff --git a/tests/test_api.py b/tests/test_api.py index e65af27..1809e2e 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -895,6 +895,82 @@ async def test_api_get_active_tasks_with_worker_filter(client: AsyncClient): assert "active_tasks" in data +@pytest.mark.asyncio +async def test_api_get_scheduled_tasks(client: AsyncClient): + """Test GET /api/workers/scheduled.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.scheduled.return_value = { + "celery@worker1": [ + {"eta": "2026-09-04T12:00:00", "priority": 3, "request": {"id": "task-1"}} + ] + } + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + response = await client.get("/api/workers/scheduled") + assert response.status_code == 200 + data = response.json() + assert "scheduled_tasks" in data + + +@pytest.mark.asyncio +async def test_api_get_scheduled_tasks_with_worker_filter(client: AsyncClient): + """Test GET /api/workers/scheduled with worker filter.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.destination.return_value = mock_inspect + mock_inspect.scheduled.return_value = { + "celery@worker1": [{"eta": "2026-09-04T12:00:00", "request": {}}] + } + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + response = await client.get("/api/workers/scheduled?worker_name=celery@worker1") + assert response.status_code == 200 + data = response.json() + assert "scheduled_tasks" in data + + +@pytest.mark.asyncio +async def test_api_get_reserved_tasks(client: AsyncClient): + """Test GET /api/workers/reserved.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.reserved.return_value = { + "celery@worker1": [{"id": "task-1", "name": "myapp.tasks.task1"}] + } + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + response = await client.get("/api/workers/reserved") + assert response.status_code == 200 + data = response.json() + assert "reserved_tasks" in data + + +@pytest.mark.asyncio +async def test_api_get_reserved_tasks_with_worker_filter(client: AsyncClient): + """Test GET /api/workers/reserved with worker filter.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.destination.return_value = mock_inspect + mock_inspect.reserved.return_value = { + "celery@worker1": [{"id": "task-1", "name": "myapp.tasks.task1"}] + } + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + response = await client.get("/api/workers/reserved?worker_name=celery@worker1") + assert response.status_code == 200 + data = response.json() + assert "reserved_tasks" in data + + @pytest.mark.asyncio async def test_api_list_queues(client: AsyncClient): """Test GET /api/queues.""" diff --git a/tests/test_workers.py b/tests/test_workers.py index 2eea8ba..c2af1ea 100644 --- a/tests/test_workers.py +++ b/tests/test_workers.py @@ -6,6 +6,8 @@ from taskowl.workers import ( get_active_tasks, + get_reserved_tasks, + get_scheduled_tasks, get_worker_stats, list_workers, scale_worker_pool, @@ -238,3 +240,137 @@ async def test_get_active_tasks_error(): assert "error" in result assert "Failed to get active tasks" in result["error"] + + +@pytest.mark.asyncio +async def test_get_scheduled_tasks_success(): + """Test getting scheduled tasks.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.scheduled.return_value = { + "celery@worker1": [ + { + "eta": "2026-09-04T12:00:00", + "priority": 3, + "request": {"id": "task-1", "name": "myapp.tasks.task1"}, + } + ] + } + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + result = await get_scheduled_tasks() + + assert "scheduled_tasks" in result + assert "celery@worker1" in result["scheduled_tasks"] + + +@pytest.mark.asyncio +async def test_get_scheduled_tasks_with_worker_filter(): + """Test getting scheduled tasks for a specific worker.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.destination.return_value = mock_inspect + mock_inspect.scheduled.return_value = { + "celery@worker1": [{"eta": "2026-09-04T12:00:00", "request": {}}] + } + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + result = await get_scheduled_tasks("celery@worker1") + + assert "scheduled_tasks" in result + mock_inspect.destination.assert_called_once_with(["celery@worker1"]) + + +@pytest.mark.asyncio +async def test_get_scheduled_tasks_no_tasks(): + """Test getting scheduled tasks when none are scheduled.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.scheduled.return_value = None + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + result = await get_scheduled_tasks() + + assert result == {"scheduled_tasks": {}} + + +@pytest.mark.asyncio +async def test_get_scheduled_tasks_error(): + """Test getting scheduled tasks when an error occurs.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_get_app.side_effect = Exception("Connection failed") + + result = await get_scheduled_tasks() + + assert "error" in result + assert "Failed to get scheduled tasks" in result["error"] + + +@pytest.mark.asyncio +async def test_get_reserved_tasks_success(): + """Test getting reserved tasks.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.reserved.return_value = { + "celery@worker1": [{"id": "task-1", "name": "myapp.tasks.task1"}] + } + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + result = await get_reserved_tasks() + + assert "reserved_tasks" in result + assert "celery@worker1" in result["reserved_tasks"] + + +@pytest.mark.asyncio +async def test_get_reserved_tasks_with_worker_filter(): + """Test getting reserved tasks for a specific worker.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.destination.return_value = mock_inspect + mock_inspect.reserved.return_value = { + "celery@worker1": [{"id": "task-1", "name": "myapp.tasks.task1"}] + } + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + result = await get_reserved_tasks("celery@worker1") + + assert "reserved_tasks" in result + mock_inspect.destination.assert_called_once_with(["celery@worker1"]) + + +@pytest.mark.asyncio +async def test_get_reserved_tasks_no_tasks(): + """Test getting reserved tasks when none are reserved.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_inspect = MagicMock() + mock_inspect.reserved.return_value = None + mock_app.control.inspect.return_value = mock_inspect + mock_get_app.return_value = mock_app + + result = await get_reserved_tasks() + + assert result == {"reserved_tasks": {}} + + +@pytest.mark.asyncio +async def test_get_reserved_tasks_error(): + """Test getting reserved tasks when an error occurs.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_get_app.side_effect = Exception("Connection failed") + + result = await get_reserved_tasks() + + assert "error" in result + assert "Failed to get reserved tasks" in result["error"]