Scheduler wait_for_task and find_task_by_name (#388)

* feat: scheduler.wait_for_task tool and context param for scheduler.run_task tool

* fix: Tasks in same context are not waitable

* feat: scheduler.find_task_by_name

* fix: find_task_by_name case insensitive

---------

Co-authored-by: Rafael Uzarowski <uzarowski.rafael@proton.me>

prompts/default/agent.system.tool.scheduler.md

@@ -11,10 +11,12 @@ dedicated_context flag then the task will run in the chat it was created in incl
 There are manual and automatically executed tasks.
 Automatic execution happens by a schedule defined when creating the task.
 
+Tasks are run asynchronously. If you need to wait for a running task's completion or need the result of the last task run, use the scheduler:wait_for_task tool. It will wait for the task completion in case the task is currently running and will provide the result of the last execution.
+
 ### Important instructions
 When a task is scheduled or planned, do not manually run it, if you have no more tasks, respond to user.
 Be careful not to create recursive prompt, do not send a message that would make the agent schedule more tasks, no need to mention the interval in message, just the objective.
-
+!!! When the user asks you to execute a task, first check if the task already exists and do not create a new task for execution. Execute the existing task instead. If the task in question does not exist ask the user what action to take. Never create tasks if asked to execute a task.
 
 ### Types of scheduler tasks
 There are 3 types of scheduler tasks:
@@ -34,7 +36,7 @@ This type of task is run manually and does not follow any schedule. It can be ru
 
 #### scheduler:list_tasks
 List all tasks present in the system with their 'uuid', 'name', 'type', 'state', 'schedule' and 'next_run'.
-All runnable tasks can be listed and filtered here. The arbuments a filter fields.
+All runnable tasks can be listed and filtered here. The arguments are filter fields.
 
 ##### Arguments:
 * state: list(str) (Optional) - The state filter, one of "idle", "running", "disabled", "error". To only show tasks in given state.
@@ -59,6 +61,26 @@ All runnable tasks can be listed and filtered here. The arbuments a filter field
 ~~~
 
 
+#### scheduler:find_task_by_name
+List all tasks whose name is matching partially or fully the provided name parameter.
+
+##### Arguments:
+* name: str - The task name to look for
+
+##### Usage:
+~~~json
+{
+    "thoughts": [
+        "I must look for tasks with name XYZ"
+    ],
+    "tool_name": "scheduler:find_task_by_name",
+    "tool_args": {
+        "name": "XYZ"
+    }
+}
+~~~
+
+
 #### scheduler:show_task
 Show task details for scheduler task with the given uuid.
 
@@ -84,9 +106,11 @@ Execute a task manually which is not in "running" state
 This can be used to trigger tasks manually.
 Normally you should only "run" tasks manually if they are in the "idle" state.
 It is also advised to only run "adhoc" tasks manually but every task type can be triggered by this tool.
+You can pass input data in text form as the "context" argument. The context will then be prepended to the task prompt when executed. This way you can pass for example result of one task as the input of another task or provide additional information specific to this one task run.
 
 ##### Arguments:
 * uuid: string - The uuid of the task to run. Can be retrieved for example from "scheduler:tasks_list"
+* context: (Optional) string - The context that will be prepended to the actual task prompt as contextual information.
 
 ##### Usage (execute task with uuid "xyz-123"):
 ~~~json
@@ -97,6 +121,7 @@ It is also advised to only run "adhoc" tasks manually but every task type can be
     "tool_name": "scheduler:run_task",
     "tool_args": {
         "uuid": "xyz-123",
+        "context": "This text is useful to execute the task more precisely"
     }
 }
 ~~~
@@ -218,3 +243,24 @@ The planned type of tasks is being run by a fixed plan, a list of datetimes that
     }
 }
 ~~~
+
+
+#### scheduler:wait_for_task
+Wait for the completion of a scheduler task identified by the uuid argument and return the result of last execution of the task.
+Attention: You can only wait for tasks running in a different chat context (dedicated). Tasks with dedicated_context=False can not be waited for.
+
+##### Arguments:
+* uuid: string - The uuid of the task to wait for. Can be retrieved for example from "scheduler:tasks_list"
+
+##### Usage (wait for task with uuid "xyz-123"):
+~~~json
+{
+    "thoughts": [
+        "I need the most current result of the task xyz-123",
+    ],
+    "tool_name": "scheduler:wait_for_task",
+    "tool_args": {
+        "uuid": "xyz-123",
+    }
+}
+~~~

python/helpers/task_scheduler.py

@@ -579,6 +579,10 @@ class SchedulerTaskList(BaseModel):
         with self._lock:
             return next((task for task in self.tasks if task.name == name), None)
 
+    def find_task_by_name(self, name: str) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]:
+        with self._lock:
+            return [task for task in self.tasks if name.lower() in task.name.lower()]
+
     async def remove_task_by_uuid(self, task_uuid: str) -> "SchedulerTaskList":
         with self._lock:
             self.tasks = [task for task in self.tasks if task.uuid != task_uuid]
@@ -622,7 +626,7 @@ class TaskScheduler:
 
     async def add_task(self, task: Union[ScheduledTask, AdHocTask, PlannedTask]) -> "TaskScheduler":
         await self._tasks.add_task(task)
-        ctx = await self._get_chat_context(task) # invoke context creation
+        ctx = await self._get_chat_context(task)  # invoke context creation
         return self
 
     async def remove_task_by_uuid(self, task_uuid: str) -> "TaskScheduler":
@@ -639,11 +643,14 @@ class TaskScheduler:
     def get_task_by_name(self, name: str) -> Union[ScheduledTask, AdHocTask, PlannedTask] | None:
         return self._tasks.get_task_by_name(name)
 
+    def find_task_by_name(self, name: str) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]:
+        return self._tasks.find_task_by_name(name)
+
     async def tick(self):
         for task in await self._tasks.get_due_tasks():
             await self._run_task(task)
 
-    async def run_task_by_uuid(self, task_uuid: str):
+    async def run_task_by_uuid(self, task_uuid: str, task_context: str | None = None):
         # First reload tasks to ensure we have the latest state
         await self._tasks.reload()
 
@@ -671,13 +678,13 @@ class TaskScheduler:
                 raise ValueError(f"Task with UUID '{task_uuid}' not found after state reset")
 
         # Run the task
-        await self._run_task(task)
+        await self._run_task(task, task_context)
 
-    async def run_task_by_name(self, name: str):
+    async def run_task_by_name(self, name: str, task_context: str | None = None):
         task = self._tasks.get_task_by_name(name)
         if task is None:
             raise ValueError(f"Task with name {name} not found")
-        await self._run_task(task)
+        await self._run_task(task, task_context)
 
     async def save(self):
         await self._tasks.save()
@@ -737,9 +744,9 @@ class TaskScheduler:
             raise ValueError(f"Context ID mismatch for task {task.name}: context {context.id} != task {task.context_id}")
         save_tmp_chat(context)
 
-    async def _run_task(self, task: Union[ScheduledTask, AdHocTask, PlannedTask]):
+    async def _run_task(self, task: Union[ScheduledTask, AdHocTask, PlannedTask], task_context: str | None = None):
 
-        async def _run_task_wrapper(task_uuid: str):
+        async def _run_task_wrapper(task_uuid: str, task_context: str | None = None):
 
             # preflight checks with a snapshot of the task
             task_snapshot: Union[ScheduledTask, AdHocTask, PlannedTask] | None = self.get_task_by_uuid(task_uuid)
@@ -799,18 +806,24 @@ class TaskScheduler:
                     for filename in attachment_filenames:
                         self._printer.print(f"- {filename}")
 
+                task_prompt = f"# Starting scheduler task '{current_task.name}' ({current_task.uuid})"
+                if task_context:
+                    task_prompt = f"## Context:\n{task_context}\n\n## Task:\n{current_task.prompt}"
+                else:
+                    task_prompt = f"## Task:\n{current_task.prompt}"
+
                 # Log the message with message_id and attachments
                 context.log.log(
                     type="user",
                     heading="User message",
-                    content=current_task.prompt,
+                    content=task_prompt,
                     kvps={"attachments": attachment_filenames},
                     id=str(uuid.uuid4()),
                 )
 
                 agent.hist_add_user_message(
                     UserMessage(
-                        message=current_task.prompt,
+                        message=task_prompt,
                         system_message=[current_task.system_prompt],
                         attachments=attachment_filenames))
 
@@ -854,7 +867,7 @@ class TaskScheduler:
                 await self._tasks.save()
 
         deferred_task = DeferredTask(thread_name=self.__class__.__name__)
-        deferred_task.start_task(_run_task_wrapper, task.uuid)
+        deferred_task.start_task(_run_task_wrapper, task.uuid, task_context)
 
         # Ensure background execution doesn't exit immediately on async await, especially in script contexts
         # This helps prevent premature exits when running from non-event-loop contexts

python/tools/scheduler.py

@@ -1,20 +1,26 @@
+import asyncio
 from datetime import datetime
 import json
 import random
 import re
 from python.helpers.tool import Tool, Response
 from python.helpers.task_scheduler import (
-    TaskScheduler, ScheduledTask, AdHocTask, PlannedTask, serialize_task, TaskState, TaskSchedule, TaskPlan, parse_datetime
+    TaskScheduler, ScheduledTask, AdHocTask, PlannedTask,
+    serialize_task, TaskState, TaskSchedule, TaskPlan, parse_datetime, serialize_datetime
 )
 from agent import AgentContext
 from python.helpers import persist_chat
 
+DEFAULT_WAIT_TIMEOUT = 300
+
 
 class SchedulerTool(Tool):
 
     async def execute(self, **kwargs):
         if self.method == "list_tasks":
             return await self.list_tasks(**kwargs)
+        elif self.method == "find_task_by_name":
+            return await self.find_task_by_name(**kwargs)
         elif self.method == "show_task":
             return await self.show_task(**kwargs)
         elif self.method == "run_task":
@@ -27,6 +33,8 @@ class SchedulerTool(Tool):
             return await self.create_adhoc_task(**kwargs)
         elif self.method == "create_planned_task":
             return await self.create_planned_task(**kwargs)
+        elif self.method == "wait_for_task":
+            return await self.wait_for_task(**kwargs)
         else:
             return Response(message=f"Unknown method '{self.name}:{self.method}'", break_loop=False)
 
@@ -51,6 +59,15 @@ class SchedulerTool(Tool):
 
         return Response(message=json.dumps(filtered_tasks, indent=4), break_loop=False)
 
+    async def find_task_by_name(self, **kwargs) -> Response:
+        name: str = kwargs.get("name", None)
+        if not name:
+            return Response(message="Task name is required", break_loop=False)
+        tasks: list[ScheduledTask | AdHocTask | PlannedTask] = TaskScheduler.get().find_task_by_name(name)
+        if not tasks:
+            return Response(message=f"Task not found: {name}", break_loop=False)
+        return Response(message=json.dumps([serialize_task(task) for task in tasks], indent=4), break_loop=False)
+
     async def show_task(self, **kwargs) -> Response:
         task_uuid: str = kwargs.get("uuid", None)
         if not task_uuid:
@@ -64,12 +81,13 @@ class SchedulerTool(Tool):
         task_uuid: str = kwargs.get("uuid", None)
         if not task_uuid:
             return Response(message="Task UUID is required", break_loop=False)
+        task_context: str | None = kwargs.get("context", None)
         task: ScheduledTask | AdHocTask | PlannedTask | None = TaskScheduler.get().get_task_by_uuid(task_uuid)
         if not task:
             return Response(message=f"Task not found: {task_uuid}", break_loop=False)
-        await TaskScheduler.get().run_task_by_uuid(task_uuid)
+        await TaskScheduler.get().run_task_by_uuid(task_uuid, task_context)
         if task.context_id == self.agent.context.id:
-            break_loop = True # break loop if task is running in the same context, otherwise it would start two conversations in one window
+            break_loop = True  # break loop if task is running in the same context, otherwise it would start two conversations in one window
         else:
             break_loop = False
         return Response(message=f"Task started: {task_uuid}", break_loop=break_loop)
@@ -199,3 +217,37 @@ class SchedulerTool(Tool):
         )
         await TaskScheduler.get().add_task(task)
         return Response(message=f"Planned task '{name}' created: {task.uuid}", break_loop=False)
+
+    async def wait_for_task(self, **kwargs) -> Response:
+        task_uuid: str = kwargs.get("uuid", None)
+        if not task_uuid:
+            return Response(message="Task UUID is required", break_loop=False)
+
+        scheduler = TaskScheduler.get()
+        task: ScheduledTask | AdHocTask | PlannedTask | None = scheduler.get_task_by_uuid(task_uuid)
+        if not task:
+            return Response(message=f"Task not found: {task_uuid}", break_loop=False)
+
+        if task.context_id == self.agent.context.id:
+            return Response(message="You can only wait for tasks running in a different chat context (dedicated_context=True).", break_loop=False)
+
+        done = False
+        elapsed = 0
+        while not done:
+            await scheduler.reload()
+            task = scheduler.get_task_by_uuid(task_uuid)
+            if not task:
+                return Response(message=f"Task not found: {task_uuid}", break_loop=False)
+
+            if task.state == TaskState.RUNNING:
+                await asyncio.sleep(1)
+                elapsed += 1
+                if elapsed > DEFAULT_WAIT_TIMEOUT:
+                    return Response(message=f"Task wait timeout ({DEFAULT_WAIT_TIMEOUT} seconds): {task_uuid}", break_loop=False)
+            else:
+                done = True
+
+        return Response(
+            message=f"*Task*: {task_uuid}\n*State*: {task.state}\n*Last run*: {serialize_datetime(task.last_run)}\n*Result*:\n{task.last_result}",
+            break_loop=False
+        )