/ gateway / session_context.py
session_context.py
  1  """
  2  Session-scoped context variables for the Hermes gateway.
  3  
  4  Replaces the previous ``os.environ``-based session state
  5  (``HERMES_SESSION_PLATFORM``, ``HERMES_SESSION_CHAT_ID``, etc.) with
  6  Python's ``contextvars.ContextVar``.
  7  
  8  **Why this matters**
  9  
 10  The gateway processes messages concurrently via ``asyncio``.  When two
 11  messages arrive at the same time the old code did:
 12  
 13      os.environ["HERMES_SESSION_THREAD_ID"] = str(context.source.thread_id)
 14  
 15  Because ``os.environ`` is *process-global*, Message A's value was
 16  silently overwritten by Message B before Message A's agent finished
 17  running.  Background-task notifications and tool calls therefore routed
 18  to the wrong thread.
 19  
 20  ``contextvars.ContextVar`` values are *task-local*: each ``asyncio``
 21  task (and any ``run_in_executor`` thread it spawns) gets its own copy,
 22  so concurrent messages never interfere.
 23  
 24  **Backward compatibility**
 25  
 26  The public helper ``get_session_env(name, default="")`` mirrors the old
 27  ``os.getenv("HERMES_SESSION_*", ...)`` calls.  Existing tool code only
 28  needs to replace the import + call site:
 29  
 30      # before
 31      import os
 32      platform = os.getenv("HERMES_SESSION_PLATFORM", "")
 33  
 34      # after
 35      from gateway.session_context import get_session_env
 36      platform = get_session_env("HERMES_SESSION_PLATFORM", "")
 37  """
 38  
 39  from contextvars import ContextVar
 40  from typing import Any
 41  
 42  # Sentinel to distinguish "never set in this context" from "explicitly set to empty".
 43  # When a contextvar holds _UNSET, we fall back to os.environ (CLI/cron compat).
 44  # When it holds "" (after clear_session_vars resets it), we return "" — no fallback.
 45  _UNSET: Any = object()
 46  
 47  # ---------------------------------------------------------------------------
 48  # Per-task session variables
 49  # ---------------------------------------------------------------------------
 50  
 51  _SESSION_PLATFORM: ContextVar = ContextVar("HERMES_SESSION_PLATFORM", default=_UNSET)
 52  _SESSION_CHAT_ID: ContextVar = ContextVar("HERMES_SESSION_CHAT_ID", default=_UNSET)
 53  _SESSION_CHAT_NAME: ContextVar = ContextVar("HERMES_SESSION_CHAT_NAME", default=_UNSET)
 54  _SESSION_THREAD_ID: ContextVar = ContextVar("HERMES_SESSION_THREAD_ID", default=_UNSET)
 55  _SESSION_USER_ID: ContextVar = ContextVar("HERMES_SESSION_USER_ID", default=_UNSET)
 56  _SESSION_USER_NAME: ContextVar = ContextVar("HERMES_SESSION_USER_NAME", default=_UNSET)
 57  _SESSION_KEY: ContextVar = ContextVar("HERMES_SESSION_KEY", default=_UNSET)
 58  
 59  # Cron auto-delivery vars — set per-job in run_job() so concurrent jobs
 60  # don't clobber each other's delivery targets.
 61  _CRON_AUTO_DELIVER_PLATFORM: ContextVar = ContextVar("HERMES_CRON_AUTO_DELIVER_PLATFORM", default=_UNSET)
 62  _CRON_AUTO_DELIVER_CHAT_ID: ContextVar = ContextVar("HERMES_CRON_AUTO_DELIVER_CHAT_ID", default=_UNSET)
 63  _CRON_AUTO_DELIVER_THREAD_ID: ContextVar = ContextVar("HERMES_CRON_AUTO_DELIVER_THREAD_ID", default=_UNSET)
 64  
 65  _VAR_MAP = {
 66      "HERMES_SESSION_PLATFORM": _SESSION_PLATFORM,
 67      "HERMES_SESSION_CHAT_ID": _SESSION_CHAT_ID,
 68      "HERMES_SESSION_CHAT_NAME": _SESSION_CHAT_NAME,
 69      "HERMES_SESSION_THREAD_ID": _SESSION_THREAD_ID,
 70      "HERMES_SESSION_USER_ID": _SESSION_USER_ID,
 71      "HERMES_SESSION_USER_NAME": _SESSION_USER_NAME,
 72      "HERMES_SESSION_KEY": _SESSION_KEY,
 73      "HERMES_CRON_AUTO_DELIVER_PLATFORM": _CRON_AUTO_DELIVER_PLATFORM,
 74      "HERMES_CRON_AUTO_DELIVER_CHAT_ID": _CRON_AUTO_DELIVER_CHAT_ID,
 75      "HERMES_CRON_AUTO_DELIVER_THREAD_ID": _CRON_AUTO_DELIVER_THREAD_ID,
 76  }
 77  
 78  
 79  def set_session_vars(
 80      platform: str = "",
 81      chat_id: str = "",
 82      chat_name: str = "",
 83      thread_id: str = "",
 84      user_id: str = "",
 85      user_name: str = "",
 86      session_key: str = "",
 87  ) -> list:
 88      """Set all session context variables and return reset tokens.
 89  
 90      Call ``clear_session_vars(tokens)`` in a ``finally`` block to restore
 91      the previous values when the handler exits.
 92  
 93      Returns a list of ``Token`` objects (one per variable) that can be
 94      passed to ``clear_session_vars``.
 95      """
 96      tokens = [
 97          _SESSION_PLATFORM.set(platform),
 98          _SESSION_CHAT_ID.set(chat_id),
 99          _SESSION_CHAT_NAME.set(chat_name),
100          _SESSION_THREAD_ID.set(thread_id),
101          _SESSION_USER_ID.set(user_id),
102          _SESSION_USER_NAME.set(user_name),
103          _SESSION_KEY.set(session_key),
104      ]
105      return tokens
106  
107  
108  def clear_session_vars(tokens: list) -> None:
109      """Mark session context variables as explicitly cleared.
110  
111      Sets all variables to ``""`` so that ``get_session_env`` returns an empty
112      string instead of falling back to (potentially stale) ``os.environ``
113      values.  The *tokens* argument is accepted for API compatibility with
114      callers that saved the return value of ``set_session_vars``, but the
115      actual clearing uses ``var.set("")`` rather than ``var.reset(token)``
116      to ensure the "explicitly cleared" state is distinguishable from
117      "never set" (which holds the ``_UNSET`` sentinel).
118      """
119      for var in (
120          _SESSION_PLATFORM,
121          _SESSION_CHAT_ID,
122          _SESSION_CHAT_NAME,
123          _SESSION_THREAD_ID,
124          _SESSION_USER_ID,
125          _SESSION_USER_NAME,
126          _SESSION_KEY,
127      ):
128          var.set("")
129  
130  
131  def get_session_env(name: str, default: str = "") -> str:
132      """Read a session context variable by its legacy ``HERMES_SESSION_*`` name.
133  
134      Drop-in replacement for ``os.getenv("HERMES_SESSION_*", default)``.
135  
136      Resolution order:
137      1. Context variable (set by the gateway for concurrency-safe access).
138         If the variable was explicitly set (even to ``""``) via
139         ``set_session_vars`` or ``clear_session_vars``, that value is
140         returned — **no fallback to os.environ**.
141      2. ``os.environ`` (only when the context variable was never set in
142         this context — i.e. CLI, cron scheduler, and test processes that
143         don't use ``set_session_vars`` at all).
144      3. *default*
145      """
146      import os
147  
148      var = _VAR_MAP.get(name)
149      if var is not None:
150          value = var.get()
151          if value is not _UNSET:
152              return value
153      # Fall back to os.environ for CLI, cron, and test compatibility
154      return os.getenv(name, default)