From 469bb2b103789379def67c67bef48387e296fd22 Mon Sep 17 00:00:00 2001
From: Javier Martinez <javiermartinezalvarez98@gmail.com>
Date: Thu, 16 Jul 2026 15:20:25 +0200
Subject: [PATCH] fix: add logs in celery runner

---
 .../celery/tasks/tools/tool_run_task.py       |  7 +-
 .../components/tools/tool_scheduler.py        | 10 ++-
 tests/components/tools/test_tool_scheduler.py | 65 ++++++++++++-------
 3 files changed, 54 insertions(+), 28 deletions(-)

diff --git a/private_gpt/celery/tasks/tools/tool_run_task.py b/private_gpt/celery/tasks/tools/tool_run_task.py
index ea3fe283..a377279d 100644
--- a/private_gpt/celery/tasks/tools/tool_run_task.py
+++ b/private_gpt/celery/tasks/tools/tool_run_task.py
@@ -26,7 +26,12 @@ logger = logging.getLogger(__name__)
     ignore_result=True,
 )
 async def tool_run_task(*, request_data: dict[str, Any]) -> None:
-    request = ToolExecutionRequest.model_validate(request_data)
+    try:
+        request = ToolExecutionRequest.model_validate(request_data)
+    except Exception:
+        logger.exception("Invalid tool execution request")
+        raise
+
     try:
         response = await execute_tool_request(
             request,
diff --git a/private_gpt/components/tools/tool_scheduler.py b/private_gpt/components/tools/tool_scheduler.py
index 07c10e86..e34f8326 100644
--- a/private_gpt/components/tools/tool_scheduler.py
+++ b/private_gpt/components/tools/tool_scheduler.py
@@ -89,9 +89,13 @@ class LocalToolScheduler(BaseToolScheduler):
         state_ctx: ChatState | None = None,
         interceptors: list[ToolExecutionInterceptor] | None = None,
     ) -> ToolExecutionResponse:
-        return await execute_tool_request(
-            request, state_ctx=state_ctx, interceptors=interceptors
-        )
+        try:
+            return await execute_tool_request(
+                request, state_ctx=state_ctx, interceptors=interceptors
+            )
+        except Exception:
+            logger.exception("Local tool '%s' execution failed", request.tool_name)
+            raise
 
     async def cancel(
         self,
diff --git a/tests/components/tools/test_tool_scheduler.py b/tests/components/tools/test_tool_scheduler.py
index 8492f844..0edda0c3 100644
--- a/tests/components/tools/test_tool_scheduler.py
+++ b/tests/components/tools/test_tool_scheduler.py
@@ -1,11 +1,48 @@
+import logging
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
 from private_gpt.components.chat.models.chat_config_models import ToolSpec
 from private_gpt.components.tools.remote_execution import ToolExecutionRequest
-from private_gpt.components.tools.tool_scheduler import CeleryToolScheduler
+from private_gpt.components.tools.tool_scheduler import (
+    CeleryToolScheduler,
+    LocalToolScheduler,
+)
+
+
+def tool_request() -> ToolExecutionRequest:
+    return ToolExecutionRequest.model_validate(
+        {
+            "tool_id": "tool-1",
+            "tool_name": "bash",
+            "tool_kwargs": {},
+            "tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}),
+            "context": {"correlation_id": "msg-1", "messages": []},
+        }
+    )
+
+
+@pytest.mark.anyio
+async def test_local_tool_scheduler_logs_execution_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    execute_tool_request = AsyncMock(side_effect=RuntimeError("boom"))
+    monkeypatch.setattr(
+        "private_gpt.components.tools.tool_scheduler.execute_tool_request",
+        execute_tool_request,
+    )
+
+    with caplog.at_level(
+        logging.ERROR,
+        logger="private_gpt.components.tools.tool_scheduler",
+    ), pytest.raises(RuntimeError, match="boom"):
+        await LocalToolScheduler().execute(tool_request())
+
+    assert "Local tool 'bash' execution failed" in caplog.text
+    assert "RuntimeError: boom" in caplog.text
 
 
 @pytest.mark.anyio
@@ -16,18 +53,8 @@ async def test_celery_tool_scheduler_execute_raises_not_implemented() -> None:
         ),
     )
 
-    request = ToolExecutionRequest.model_validate(
-        {
-            "tool_id": "tool-1",
-            "tool_name": "bash",
-            "tool_kwargs": {},
-            "tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}),
-            "context": {"correlation_id": "msg-1", "messages": []},
-        }
-    )
-
     with pytest.raises(NotImplementedError):
-        await scheduler.execute(request)
+        await scheduler.execute(tool_request())
 
 
 @pytest.mark.anyio
@@ -43,17 +70,7 @@ async def test_celery_tool_scheduler_cancel_revokes_task(
         ),
     )
 
-    request = ToolExecutionRequest.model_validate(
-        {
-            "tool_id": "tool-1",
-            "tool_name": "bash",
-            "tool_kwargs": {},
-            "tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}),
-            "context": {"correlation_id": "msg-1", "messages": []},
-        }
-    )
-
-    cancelled = await scheduler.cancel(request, task_id="task-abc")
+    cancelled = await scheduler.cancel(tool_request(), task_id="task-abc")
 
     assert cancelled is True
     celery_app.control.revoke.assert_called_once_with("task-abc", terminate=True)