Merge PR #1458: prevent cron self-scheduling safely
This commit is contained in:
@@ -1,5 +1,6 @@
|
|||||||
"""Cron tool for scheduling reminders and tasks."""
|
"""Cron tool for scheduling reminders and tasks."""
|
||||||
|
|
||||||
|
from contextvars import ContextVar
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from nanobot.agent.tools.base import Tool
|
from nanobot.agent.tools.base import Tool
|
||||||
@@ -14,12 +15,21 @@ class CronTool(Tool):
|
|||||||
self._cron = cron_service
|
self._cron = cron_service
|
||||||
self._channel = ""
|
self._channel = ""
|
||||||
self._chat_id = ""
|
self._chat_id = ""
|
||||||
|
self._in_cron_context: ContextVar[bool] = ContextVar("cron_in_context", default=False)
|
||||||
|
|
||||||
def set_context(self, channel: str, chat_id: str) -> None:
|
def set_context(self, channel: str, chat_id: str) -> None:
|
||||||
"""Set the current session context for delivery."""
|
"""Set the current session context for delivery."""
|
||||||
self._channel = channel
|
self._channel = channel
|
||||||
self._chat_id = chat_id
|
self._chat_id = chat_id
|
||||||
|
|
||||||
|
def set_cron_context(self, active: bool):
|
||||||
|
"""Mark whether the tool is executing inside a cron job callback."""
|
||||||
|
return self._in_cron_context.set(active)
|
||||||
|
|
||||||
|
def reset_cron_context(self, token) -> None:
|
||||||
|
"""Restore previous cron context."""
|
||||||
|
self._in_cron_context.reset(token)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def name(self) -> str:
|
def name(self) -> str:
|
||||||
return "cron"
|
return "cron"
|
||||||
@@ -72,6 +82,8 @@ class CronTool(Tool):
|
|||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> str:
|
) -> str:
|
||||||
if action == "add":
|
if action == "add":
|
||||||
|
if self._in_cron_context.get():
|
||||||
|
return "Error: cannot schedule new jobs from within a cron job execution"
|
||||||
return self._add_job(message, every_seconds, cron_expr, tz, at)
|
return self._add_job(message, every_seconds, cron_expr, tz, at)
|
||||||
elif action == "list":
|
elif action == "list":
|
||||||
return self._list_jobs()
|
return self._list_jobs()
|
||||||
|
|||||||
@@ -296,6 +296,7 @@ def gateway(
|
|||||||
# Set cron callback (needs agent)
|
# Set cron callback (needs agent)
|
||||||
async def on_cron_job(job: CronJob) -> str | None:
|
async def on_cron_job(job: CronJob) -> str | None:
|
||||||
"""Execute a cron job through the agent."""
|
"""Execute a cron job through the agent."""
|
||||||
|
from nanobot.agent.tools.cron import CronTool
|
||||||
from nanobot.agent.tools.message import MessageTool
|
from nanobot.agent.tools.message import MessageTool
|
||||||
reminder_note = (
|
reminder_note = (
|
||||||
"[Scheduled Task] Timer finished.\n\n"
|
"[Scheduled Task] Timer finished.\n\n"
|
||||||
@@ -303,12 +304,21 @@ def gateway(
|
|||||||
f"Scheduled instruction: {job.payload.message}"
|
f"Scheduled instruction: {job.payload.message}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Prevent the agent from scheduling new cron jobs during execution
|
||||||
|
cron_tool = agent.tools.get("cron")
|
||||||
|
cron_token = None
|
||||||
|
if isinstance(cron_tool, CronTool):
|
||||||
|
cron_token = cron_tool.set_cron_context(True)
|
||||||
|
try:
|
||||||
response = await agent.process_direct(
|
response = await agent.process_direct(
|
||||||
reminder_note,
|
reminder_note,
|
||||||
session_key=f"cron:{job.id}",
|
session_key=f"cron:{job.id}",
|
||||||
channel=job.payload.channel or "cli",
|
channel=job.payload.channel or "cli",
|
||||||
chat_id=job.payload.to or "direct",
|
chat_id=job.payload.to or "direct",
|
||||||
)
|
)
|
||||||
|
finally:
|
||||||
|
if isinstance(cron_tool, CronTool) and cron_token is not None:
|
||||||
|
cron_tool.reset_cron_context(cron_token)
|
||||||
|
|
||||||
message_tool = agent.tools.get("message")
|
message_tool = agent.tools.get("message")
|
||||||
if isinstance(message_tool, MessageTool) and message_tool._sent_in_turn:
|
if isinstance(message_tool, MessageTool) and message_tool._sent_in_turn:
|
||||||
|
|||||||
Reference in New Issue
Block a user