Worker pool restart

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`, `get_scheduled_tasks`, `get_reserved_tasks` |
+| **Workers** | `get_worker_status`, `list_workers`, `get_worker_stats`, `shutdown_worker`, `scale_worker_pool`, `restart_worker_pool`, `get_active_tasks`, `get_scheduled_tasks`, `get_reserved_tasks` |
 | **Queues** | `list_queues` |
 
-**Total: 19 tools**
+**Total: 20 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` |
+| "Restart the pool on celery@worker1" | `restart_worker_pool` |
 | "What's scheduled to run next?" | `get_scheduled_tasks`, `get_reserved_tasks` |
 | "Retry task abc" | `retry_task` |
 | "Run myapp.tasks.process now" | `execute_task` |
@@ -264,7 +265,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`, `GET /api/workers/scheduled`, `GET /api/workers/reserved` |
-| **Worker actions** | `POST /api/workers/{name}/shutdown`, `POST /api/workers/{name}/scale` |
+| **Worker actions** | `POST /api/workers/{name}/shutdown`, `POST /api/workers/{name}/scale`, `POST /api/workers/{name}/restart` |
 | **Queues** | `GET /api/queues` |
 | **Ops** | `GET /health`, `GET /metrics` |
 

src/taskowl/main.py

@@ -31,6 +31,7 @@
     get_scheduled_tasks,
     get_worker_stats,
     list_workers,
+    restart_worker_pool,
     scale_worker_pool,
     shutdown_worker,
 )
@@ -342,6 +343,24 @@ async def api_scale_worker_pool(
     return result
 
 
+@app.post("/api/workers/{worker_name}/restart")
+async def api_restart_worker_pool(
+    worker_name: str,
+    reload: bool = False,
+    _: None = Depends(verify_api_key),
+) -> dict:
+    """Restart a worker's execution pool.
+
+    Args:
+        worker_name: Name of the worker
+        reload: If True, reload modules when restarting the pool
+    """
+    result = await restart_worker_pool(worker_name, reload)
+    if "error" in result:
+        raise HTTPException(status_code=400, detail=result["error"])
+    return result
+
+
 @app.get("/api/workers/active-tasks")
 async def api_get_active_tasks(
     worker_name: str | None = None,

src/taskowl/mcp/tools.py

@@ -356,6 +356,26 @@ async def scale_worker_pool(worker_name: str, delta: int) -> dict:
             response.raise_for_status()
             return response.json()
 
+    @server.tool(
+        name="restart_worker_pool",
+        description="Restart a worker's execution pool (optionally reloading modules)",
+    )
+    async def restart_worker_pool(worker_name: str, reload: bool = False) -> dict:
+        """Restart a worker's execution pool.
+
+        Args:
+            worker_name: Name of the worker
+            reload: If True, reload modules when restarting the pool
+        """
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"http://{settings.taskowl_host}:{settings.taskowl_port}/api/workers/{worker_name}/restart",
+                params={"reload": reload},
+                headers=_get_headers(),
+            )
+            response.raise_for_status()
+            return response.json()
+
     @server.tool(
         name="get_active_tasks",
         description="Get currently executing tasks across all workers or a specific worker",

src/taskowl/workers.py

@@ -124,6 +124,29 @@ async def scale_worker_pool(worker_name: str, delta: int) -> dict:
         return {"error": f"Failed to scale worker pool: {str(e)}"}
 
 
+async def restart_worker_pool(worker_name: str, reload: bool = False) -> dict:
+    """Restart a worker's execution pool.
+
+    Args:
+        worker_name: Name of the worker
+        reload: If True, reload modules when restarting the pool
+
+    Returns:
+        Dict with status and message
+    """
+    try:
+        app = _get_celery_app()
+        app.control.pool_restart(destination=[worker_name], reload=reload)
+
+        return {
+            "status": "success",
+            "message": f"Pool restart command sent to {worker_name}",
+            "worker": worker_name,
+        }
+    except Exception as e:
+        return {"error": f"Failed to restart worker pool: {str(e)}"}
+
+
 async def get_active_tasks(worker_name: str | None = None) -> dict:
     """Get currently executing tasks.
 

tests/test_api.py

@@ -858,6 +858,52 @@ async def test_api_scale_worker_pool_zero_delta(client: AsyncClient):
     assert "detail" in data
 
 
+@pytest.mark.asyncio
+async def test_api_restart_worker_pool(client: AsyncClient):
+    """Test POST /api/workers/{worker_name}/restart."""
+    with patch("taskowl.workers._get_celery_app") as mock_get_app:
+        mock_app = MagicMock()
+        mock_get_app.return_value = mock_app
+
+        response = await client.post("/api/workers/celery@worker1/restart")
+        assert response.status_code == 200
+        data = response.json()
+        assert data["status"] == "success"
+        assert data["worker"] == "celery@worker1"
+        mock_app.control.pool_restart.assert_called_once_with(
+            destination=["celery@worker1"], reload=False
+        )
+
+
+@pytest.mark.asyncio
+async def test_api_restart_worker_pool_reload(client: AsyncClient):
+    """Test POST /api/workers/{worker_name}/restart with reload=true."""
+    with patch("taskowl.workers._get_celery_app") as mock_get_app:
+        mock_app = MagicMock()
+        mock_get_app.return_value = mock_app
+
+        response = await client.post("/api/workers/celery@worker1/restart?reload=true")
+        assert response.status_code == 200
+        data = response.json()
+        assert data["status"] == "success"
+        mock_app.control.pool_restart.assert_called_once_with(
+            destination=["celery@worker1"], reload=True
+        )
+
+
+@pytest.mark.asyncio
+async def test_api_restart_worker_pool_error(client: AsyncClient):
+    """Test POST /api/workers/{worker_name}/restart on error."""
+    with patch("taskowl.workers._get_celery_app") as mock_get_app:
+        mock_app = MagicMock()
+        mock_app.control.pool_restart.side_effect = Exception("Restart failed")
+        mock_get_app.return_value = mock_app
+
+        response = await client.post("/api/workers/celery@worker1/restart")
+        assert response.status_code == 400
+        assert "Restart failed" in response.json()["detail"]
+
+
 @pytest.mark.asyncio
 async def test_api_get_active_tasks(client: AsyncClient):
     """Test GET /api/workers/active-tasks."""

tests/test_workers.py

@@ -10,6 +10,7 @@
     get_scheduled_tasks,
     get_worker_stats,
     list_workers,
+    restart_worker_pool,
     scale_worker_pool,
     shutdown_worker,
 )
@@ -178,6 +179,51 @@ async def test_scale_worker_pool_error():
         assert "Failed to scale worker pool" in result["error"]
 
 
+@pytest.mark.asyncio
+async def test_restart_worker_pool_success():
+    """Test restarting a worker pool."""
+    with patch("taskowl.workers._get_celery_app") as mock_get_app:
+        mock_app = MagicMock()
+        mock_get_app.return_value = mock_app
+
+        result = await restart_worker_pool("celery@worker1")
+
+        assert result["status"] == "success"
+        assert result["worker"] == "celery@worker1"
+        mock_app.control.pool_restart.assert_called_once_with(
+            destination=["celery@worker1"], reload=False
+        )
+
+
+@pytest.mark.asyncio
+async def test_restart_worker_pool_reload():
+    """Test restarting a worker pool with reload enabled."""
+    with patch("taskowl.workers._get_celery_app") as mock_get_app:
+        mock_app = MagicMock()
+        mock_get_app.return_value = mock_app
+
+        result = await restart_worker_pool("celery@worker1", reload=True)
+
+        assert result["status"] == "success"
+        mock_app.control.pool_restart.assert_called_once_with(
+            destination=["celery@worker1"], reload=True
+        )
+
+
+@pytest.mark.asyncio
+async def test_restart_worker_pool_error():
+    """Test restarting a worker pool when an error occurs."""
+    with patch("taskowl.workers._get_celery_app") as mock_get_app:
+        mock_app = MagicMock()
+        mock_app.control.pool_restart.side_effect = Exception("Restart failed")
+        mock_get_app.return_value = mock_app
+
+        result = await restart_worker_pool("celery@worker1")
+
+        assert "error" in result
+        assert "Failed to restart worker pool" in result["error"]
+
+
 @pytest.mark.asyncio
 async def test_get_active_tasks_success():
     """Test getting active tasks."""