diff --git a/api/streaming.py b/api/streaming.py index b46ca813d..f86bf49c0 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -4,6 +4,7 @@ Includes Sprint 10 cancel support via CANCEL_FLAGS. """ import base64 import contextlib +import contextvars import json import logging import mimetypes @@ -85,6 +86,78 @@ _ENV_LOCK = threading.Lock() _KEYLESS_CUSTOM_API_KEY = "dummy-key" +_STREAMING_CRON_PROFILE_HOME: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "webui_streaming_cron_profile_home", + default=None, +) +_STREAMING_CRONJOB_WRAPPER_INSTALLED = False + + +def _install_streaming_cronjob_profile_wrapper() -> None: + """Wrap the agent cronjob tool so calls run under the streaming profile. + + The in-chat agent run already binds per-turn contextvars for other + session-scoped state. Cron jobs are special because ``cron.jobs`` snapshots + path constants at import time, so the model-facing ``cronjob`` tool must + enter the existing WebUI cron profile context at the tool-call boundary. + That context uses the cron-specific lock and restores the module caches as + soon as the single cron tool call returns, avoiding long-lived global path + mutation for the whole agent turn. + """ + global _STREAMING_CRONJOB_WRAPPER_INSTALLED + if _STREAMING_CRONJOB_WRAPPER_INSTALLED: + return + try: + from tools.registry import registry + except Exception: + logger.debug("streaming cronjob wrapper: tools registry unavailable", exc_info=True) + return + + entry = registry.get_entry("cronjob") + if entry is None: + try: + import tools.cronjob_tools # noqa: F401 + except Exception: + logger.debug("streaming cronjob wrapper: cronjob tool import failed", exc_info=True) + entry = registry.get_entry("cronjob") + if entry is None: + logger.debug("streaming cronjob wrapper: cronjob tool not registered") + return + original_handler = entry.handler + if getattr(original_handler, "_webui_streaming_profile_wrapper", False): + _STREAMING_CRONJOB_WRAPPER_INSTALLED = True + return + + # This relies on the agent tool executor's ``propagate_context_to_thread`` + # using an unfiltered ``contextvars.copy_context()`` so this WebUI-owned + # contextvar reaches the sync cronjob handler even when the tool call runs + # on the agent's ThreadPoolExecutor worker. + def _profile_scoped_cronjob_handler(args, **kwargs): + profile_home = _STREAMING_CRON_PROFILE_HOME.get() + if not profile_home: + return original_handler(args, **kwargs) + from api.profiles import cron_profile_context_for_home + with cron_profile_context_for_home(Path(profile_home)): + return original_handler(args, **kwargs) + + _profile_scoped_cronjob_handler.__dict__["_webui_streaming_profile_wrapper"] = True + _profile_scoped_cronjob_handler.__dict__["_webui_original_handler"] = original_handler + registry.register( + name=entry.name, + toolset=entry.toolset, + schema=entry.schema, + handler=_profile_scoped_cronjob_handler, + check_fn=entry.check_fn, + requires_env=entry.requires_env, + is_async=entry.is_async, + description=entry.description, + emoji=entry.emoji, + max_result_size_chars=entry.max_result_size_chars, + dynamic_schema_overrides=entry.dynamic_schema_overrides, + ) + _STREAMING_CRONJOB_WRAPPER_INSTALLED = True + + _PERSISTENT_MEMORY_FILES = ( ("memory", ("memories", "MEMORY.md")), ("user", ("memories", "USER.md")), @@ -5981,6 +6054,7 @@ def _run_agent_streaming( # Placed ABOVE the _checkpoint_stop cluster so that cluster stays adjacent # to the `try:` (preserves the Issue #765 static-locator invariant). _turn_session_identity_tokens = None + _streaming_cron_profile_home_token = None # Initialised here (before any code that may raise) so the outer `finally` # block can safely check `if _checkpoint_stop is not None` even when an # exception fires before the checkpoint thread is created (Issue #765). @@ -6027,6 +6101,7 @@ def _run_agent_streaming( ) _profile_home_path = get_hermes_home_for_profile(getattr(s, 'profile', None)) _profile_home = str(_profile_home_path) + _streaming_cron_profile_home_token = _STREAMING_CRON_PROFILE_HOME.set(_profile_home) _profile_runtime_env = get_profile_runtime_env(_profile_home_path) _safe_profile_runtime_env = filter_runtime_env_for_gateway_parity(_profile_runtime_env) except ImportError: @@ -6084,6 +6159,7 @@ def _run_agent_streaming( # first-time module initialisation (which can be slow) does not # block other concurrent sessions waiting on _ENV_LOCK (#2024). _prewarm_skill_tool_modules() + _install_streaming_cronjob_profile_wrapper() # Still set process-level env as fallback for tools that bypass thread-local # Acquire lock only for the env mutation, then release before the agent runs. # The finally block re-acquires to restore — keeping critical sections short @@ -6108,9 +6184,12 @@ def _run_agent_streaming( os.environ['HERMES_SESSION_CHAT_ID'] = str(session_id) if _profile_home: os.environ['HERMES_HOME'] = _profile_home - # Patch module-level caches to match the active profile. + # Patch skill module caches to match the active profile. # _set_hermes_home() does this for process-wide switches - # but per-request switches skip it (#1700). + # but per-request switches skip it (#1700). The in-chat + # cronjob tool is wrapped separately at its tool-call boundary + # with cron_profile_context_for_home (#4580) so cron.jobs path + # caches are not mutated for the entire agent turn. # Modules were prewarmed by _prewarm_skill_tool_modules() # above, so we only do lightweight sys.modules lookups and # attribute assignments here — no first-time import under @@ -8770,6 +8849,8 @@ def _run_agent_streaming( update_active_run(stream_id, phase="finalizing") _last_resort_sync_from_core(s, stream_id, _agent_lock) _clear_thread_env() # TD1: always clear thread-local context + if _streaming_cron_profile_home_token is not None: + _STREAMING_CRON_PROFILE_HOME.reset(_streaming_cron_profile_home_token) # xsession wakeup misroute root fix (Option 1): restore the per-turn # session-identity context-locals (reset-token semantics). MUST run on # every exit path so a reused thread-pool worker leaks no identity and diff --git a/tests/test_scheduled_jobs_profile_isolation.py b/tests/test_scheduled_jobs_profile_isolation.py index 7d7a2f80b..6343a7aeb 100644 --- a/tests/test_scheduled_jobs_profile_isolation.py +++ b/tests/test_scheduled_jobs_profile_isolation.py @@ -323,3 +323,229 @@ def test_cron_worker_does_not_silently_fall_back_on_profile_context_failure(): "cron subprocess target appears to catch profile-context setup before " "entering the context; do not fall back to an unpinned run_job call." ) + + +def test_streaming_cronjob_wrapper_uses_profile_context_only_for_tool_call(tmp_path, monkeypatch): + """The chat cronjob fix must use the cron profile context at call time. + + Holding cron.jobs module globals for an entire streaming run would race with + other profiles. The wrapper should instead enter the existing locked cron + context only while the cronjob tool handler itself executes. + """ + import types + + from api import profiles as p + from api import streaming as st + + profile_home = tmp_path / "home" / "profiles" / "ops" + events = [] + + class Ctx: + def __init__(self, home): + self.home = str(home) + + def __enter__(self): + events.append(("enter", self.home)) + return self + + def __exit__(self, exc_type, exc, tb): + events.append(("exit", self.home)) + return False + + def original_handler(args, **kwargs): + events.append(("handler", args.get("action"), st._STREAMING_CRON_PROFILE_HOME.get())) + return "ok" + + class Entry: + name = "cronjob" + toolset = "cronjob" + schema = {"name": "cronjob"} + handler = staticmethod(original_handler) + check_fn = None + requires_env = [] + is_async = False + description = "" + emoji = "⏰" + max_result_size_chars = None + dynamic_schema_overrides = None + + entry = Entry() + + class Registry: + def get_entry(self, name): + assert name == "cronjob" + return entry + + def register(self, **kwargs): + events.append(("register", kwargs["name"])) + entry.handler = kwargs["handler"] + + tools_pkg = types.ModuleType("tools") + registry_mod = types.ModuleType("tools.registry") + registry_mod.__dict__["registry"] = Registry() + monkeypatch.setitem(sys.modules, "tools", tools_pkg) + monkeypatch.setitem(sys.modules, "tools.registry", registry_mod) + monkeypatch.setattr(st, "_STREAMING_CRONJOB_WRAPPER_INSTALLED", False) + monkeypatch.setattr(p, "cron_profile_context_for_home", Ctx) + + st._install_streaming_cronjob_profile_wrapper() + assert events == [("register", "cronjob")] + + token = st._STREAMING_CRON_PROFILE_HOME.set(str(profile_home)) + try: + assert entry.handler({"action": "list"}, task_id="t1") == "ok" + finally: + st._STREAMING_CRON_PROFILE_HOME.reset(token) + + assert events == [ + ("register", "cronjob"), + ("enter", str(profile_home)), + ("handler", "list", str(profile_home)), + ("exit", str(profile_home)), + ] + + +def test_streaming_cronjob_wrapper_context_survives_threadpool_context_copy(tmp_path, monkeypatch): + """Lock the cross-thread contextvar contract used by agent tool dispatch. + + WebUI sets the active profile on the streaming thread, then the Hermes agent + dispatches sync tools on a ThreadPoolExecutor worker under a copied + contextvars context. The cronjob wrapper must still see the profile context + on that worker or it silently falls back to the default profile. + """ + import concurrent.futures + import contextvars + import types + + from api import profiles as p + from api import streaming as st + + profile_home = tmp_path / "home" / "profiles" / "ops" + events = [] + + class Ctx: + def __init__(self, home): + self.home = str(home) + + def __enter__(self): + events.append(("enter", self.home)) + return self + + def __exit__(self, exc_type, exc, tb): + events.append(("exit", self.home)) + return False + + def original_handler(args, **kwargs): + events.append(("handler", args.get("action"), st._STREAMING_CRON_PROFILE_HOME.get())) + return "ok" + + class Entry: + name = "cronjob" + toolset = "cronjob" + schema = {"name": "cronjob"} + handler = staticmethod(original_handler) + check_fn = None + requires_env = [] + is_async = False + description = "" + emoji = "⏰" + max_result_size_chars = None + dynamic_schema_overrides = None + + entry = Entry() + + class Registry: + def get_entry(self, name): + assert name == "cronjob" + return entry + + def register(self, **kwargs): + events.append(("register", kwargs["name"])) + entry.handler = kwargs["handler"] + + registry_mod = types.ModuleType("tools.registry") + registry_mod.__dict__["registry"] = Registry() + monkeypatch.setitem(sys.modules, "tools.registry", registry_mod) + monkeypatch.setattr(st, "_STREAMING_CRONJOB_WRAPPER_INSTALLED", False) + monkeypatch.setattr(p, "cron_profile_context_for_home", Ctx) + + st._install_streaming_cronjob_profile_wrapper() + token = st._STREAMING_CRON_PROFILE_HOME.set(str(profile_home)) + try: + copied_context = contextvars.copy_context() + finally: + st._STREAMING_CRON_PROFILE_HOME.reset(token) + + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(copied_context.run, entry.handler, {"action": "list"}) + assert future.result(timeout=5) == "ok" + + assert events == [ + ("register", "cronjob"), + ("enter", str(profile_home)), + ("handler", "list", str(profile_home)), + ("exit", str(profile_home)), + ] + + +def test_streaming_cronjob_wrapper_leaves_calls_unwrapped_without_streaming_profile(monkeypatch): + """CLI/default calls through the registered cronjob handler are unchanged.""" + import types + + from api import streaming as st + + events = [] + + def original_handler(args, **kwargs): + events.append(("handler", args.get("action"))) + return "ok" + + class Entry: + name = "cronjob" + toolset = "cronjob" + schema = {"name": "cronjob"} + handler = staticmethod(original_handler) + check_fn = None + requires_env = [] + is_async = False + description = "" + emoji = "⏰" + max_result_size_chars = None + dynamic_schema_overrides = None + + entry = Entry() + + class Registry: + def get_entry(self, name): + assert name == "cronjob" + return entry + + def register(self, **kwargs): + entry.handler = kwargs["handler"] + + registry_mod = types.ModuleType("tools.registry") + registry_mod.__dict__["registry"] = Registry() + monkeypatch.setitem(sys.modules, "tools.registry", registry_mod) + monkeypatch.setattr(st, "_STREAMING_CRONJOB_WRAPPER_INSTALLED", False) + + st._install_streaming_cronjob_profile_wrapper() + + assert entry.handler({"action": "list"}) == "ok" + assert events == [("handler", "list")] + + +def test_streaming_profile_home_mutation_avoids_long_lived_cron_cache_patch(): + """Guard the streaming integration seam for issue #4580. + + Streaming still mutates process env briefly for legacy fallbacks, but cron + path caches must be scoped to the cronjob tool-call boundary via the + wrapper/contextvar path — not patched for the full agent turn. + """ + from pathlib import Path + + src = (Path(__file__).resolve().parent.parent / "api" / "streaming.py").read_text(encoding="utf-8") + assert "_install_streaming_cronjob_profile_wrapper()" in src + assert "_STREAMING_CRON_PROFILE_HOME.set(_profile_home)" in src + assert "_STREAMING_CRON_PROFILE_HOME.reset(_streaming_cron_profile_home_token)" in src + assert "def _patch_streaming_profile_module_caches" not in src + assert "old_profile_module_cache_snapshot" not in src