fix(cron): isolate cron-execution guard with contextvars

This commit is contained in:
Re-bin
2026-03-03 05:36:48 +00:00
parent ec6430fa0c
commit 30803afec0
2 changed files with 13 additions and 7 deletions

View File

@@ -1,5 +1,6 @@
"""Cron tool for scheduling reminders and tasks."""
from contextvars import ContextVar
from typing import Any
from nanobot.agent.tools.base import Tool
@@ -14,16 +15,20 @@ class CronTool(Tool):
self._cron = cron_service
self._channel = ""
self._chat_id = ""
self._in_cron_context = False
self._in_cron_context: ContextVar[bool] = ContextVar("cron_in_context", default=False)
def set_context(self, channel: str, chat_id: str) -> None:
"""Set the current session context for delivery."""
self._channel = channel
self._chat_id = chat_id
def set_cron_context(self, active: bool) -> None:
def set_cron_context(self, active: bool):
"""Mark whether the tool is executing inside a cron job callback."""
self._in_cron_context = active
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
def name(self) -> str:
@@ -77,7 +82,7 @@ class CronTool(Tool):
**kwargs: Any,
) -> str:
if action == "add":
if self._in_cron_context:
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)
elif action == "list":

View File

@@ -306,8 +306,9 @@ def gateway(
# 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_tool.set_cron_context(True)
cron_token = cron_tool.set_cron_context(True)
try:
response = await agent.process_direct(
reminder_note,
@@ -316,8 +317,8 @@ def gateway(
chat_id=job.payload.to or "direct",
)
finally:
if isinstance(cron_tool, CronTool):
cron_tool.set_cron_context(False)
if isinstance(cron_tool, CronTool) and cron_token is not None:
cron_tool.reset_cron_context(cron_token)
message_tool = agent.tools.get("message")
if isinstance(message_tool, MessageTool) and message_tool._sent_in_turn: