fix: GlobalRuntime direct access for auxiliary memory

ak-py/src/agentkernel/api/a2a/a2a.py

@@ -9,8 +9,9 @@
 from a2a.utils import new_agent_text_message
 from a2a.utils.errors import ServerError
 
-from ...core import Agent, AgentService, GlobalRuntime
+from ...core import Agent, AgentService
 from ...core.config import AKConfig
+from ...core.runtime import Runtime
 
 
 class A2A:
@@ -69,7 +70,7 @@ def _build(cls):
             return
         if not AKConfig.get().a2a.enabled:
             return
-        agents: dict[str, Agent] = GlobalRuntime.instance().agents()
+        agents: dict[str, Agent] = Runtime.current().agents()
         for name, agent in agents.items():
             whitelisted = AKConfig.get().a2a.agents == ["*"] or name in AKConfig.get().a2a.agents
             if not whitelisted:

ak-py/src/agentkernel/api/handler.py

@@ -17,7 +17,8 @@
     AgentRequestText,
 )
 
-from ..core import AgentService, Config, GlobalRuntime
+from ..core import AgentService, Config
+from ..core.runtime import Runtime
 
 
 class RESTRequestHandler(ABC):
@@ -37,7 +38,8 @@ def health():
 
         @router.get("/agents")
         def list_agents():
-            return {"agents": list(GlobalRuntime.instance().agents().keys())}
+            from ..core.runtime import Runtime
+            return {"agents": list(Runtime.current().agents().keys())}
 
         """
         pass
@@ -93,7 +95,7 @@ def health():
 
         @router.get("/agents")
         def list_agents():
-            return {"agents": list(GlobalRuntime.instance().agents().keys())}
+            return {"agents": list(Runtime.current().agents().keys())}
 
         @router.post("/run")
         async def run(body: AgentRESTRequestHandler.RunRequest):

ak-py/src/agentkernel/api/mcp/akmcp.py

@@ -4,8 +4,9 @@
 from fastmcp import Context, FastMCP
 from fastmcp.server.http import StarletteWithLifespan
 
-from ...core import Agent, AgentService, GlobalRuntime
+from ...core import Agent, AgentService
 from ...core.config import AKConfig
+from ...core.runtime import Runtime
 
 
 class MCP:
@@ -60,7 +61,7 @@ def _build(cls):
         if cls._fastmcp is None:
             cls._fastmcp = FastMCP("Agent Kernel FastMCP Instance")
         if AKConfig.get().mcp.expose_agents:
-            agents: dict[str, Agent] = GlobalRuntime.instance().agents()
+            agents: dict[str, Agent] = Runtime.current().agents()
             for name, agent in agents.items():
                 whitelisted = AKConfig.get().mcp.agents == ["*"] or name in AKConfig.get().mcp.agents
                 if not whitelisted:

ak-py/src/agentkernel/core/__init__.py

@@ -24,7 +24,7 @@
 )
 from .config import AKConfig as Config
 from .module import Module
-from .runtime import GlobalRuntime, Runtime
+from .runtime import Runtime, AuxiliaryCache
 from .service import AgentService
 from .hooks import PreHook, PostHook
 from .util.key_value_cache import KeyValueCache

ak-py/src/agentkernel/core/module.py

@@ -3,7 +3,7 @@
 
 from ..core.hooks import PostHook, PreHook
 from .base import Agent
-from .runtime import ModuleLoader
+from .runtime import Runtime
 
 
 class Module(ABC):
@@ -33,7 +33,7 @@ def unload(self):
         Unloads and deregisters all agents in the module
         """
         for agent in self._agents:
-            ModuleLoader.runtime().deregister(agent)
+            Runtime.current().deregister(agent)
         self._agents.clear()
 
     def get_agent(self, name: str) -> Agent | None:
@@ -67,7 +67,7 @@ def load(self, agents: list[Any]) -> "Module":
         for agent in agents:
             try:
                 wrapped = self._wrap(agent, agents)
-                ModuleLoader.runtime().register(wrapped)
+                Runtime.current().register(wrapped)
                 registered.append(wrapped)
             except Exception:
                 self._agents = registered

ak-py/src/agentkernel/core/runtime.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import importlib
 import logging
 from threading import RLock
@@ -6,8 +8,7 @@
 
 from singleton_type import Singleton
 
-from agentkernel.core.util.key_value_cache import KeyValueCache
-
+from ..core.util.key_value_cache import KeyValueCache
 from .base import Agent, Session
 from .builder import SessionStoreBuilder
 from .model import (
@@ -28,6 +29,9 @@ class Runtime:
     Runtime class provides the environment for hosting and running agents.
     """
 
+    _current: Optional[Runtime] = None
+    _lock: RLock = RLock()
+
     def __init__(self, sessions: SessionStore):
         """
         Initialize the Runtime.
@@ -38,36 +42,55 @@ def __init__(self, sessions: SessionStore):
         self._agents = {}
         self._sessions = sessions
 
+    @staticmethod
+    def current() -> Runtime:
+        """
+        Return the currently active Runtime instance. By default this is the
+        global singleton Runtime instance.
+
+        :return: The currently active runtime instance.
+        """
+        return Runtime._current or GlobalRuntime.instance()
+
     def __enter__(self) -> "Runtime":
         """
-        Enter the Runtime context manager and attach the Runtime to the ModuleLoader.
+        Enter the Runtime context manager and set as the current Runtime.
 
-        This method is called when entering a 'with' statement block. It attaches
-        the ModuleLoader to this runtime instance, making it the active runtime
-        context for module loading operations.
+        This method is called when entering a 'with' statement block. It sets
+        this runtime instance as the active runtime context.
 
         :return: The runtime instance itself, allowing it to be used as a context manager in with statements.
         """
-        ModuleLoader.attach(self)
+        with Runtime._lock:
+            if Runtime._current is not None and Runtime._current != self:
+                raise Exception("A different runtime is already active")
+            Runtime._current = self
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         """
-        Exit the Runtime context manager and detach from the ModuleLoader.
+        Exit the Runtime context manager and clear the current Runtime.
 
-        This method is called when exiting a 'with' statement block. It detaches
-        the runtime instance from the ModuleLoader, performing necessary cleanup.
+        This method is called when exiting a 'with' statement block. It clears
+        the runtime instance from being the active runtime, performing necessary cleanup.
         """
-        ModuleLoader.detach(self)
+        with Runtime._lock:
+            if Runtime._current is not None and Runtime._current != self:
+                raise Exception("A different runtime is currently active")
+            Runtime._current = None
 
     def load(self, module: str) -> ModuleType:
         """
         Loads an agent module dynamically.
         :param module: Name of the module to load.
         :return: The loaded module.
+
+        :raises ModuleNotFoundError: If the specified module cannot be found.
+        :raises ImportError: If there's an error during the module import process.
         """
         self._log.debug(f"Loading module '{module}'")
-        return ModuleLoader.load(self, module)
+        with self:
+            return importlib.import_module(module)
 
     def agents(self) -> dict[str, Agent]:
         """
@@ -154,32 +177,6 @@ def sessions(self) -> SessionStore:
         """
         return self._sessions
 
-    def get_volatile_cache(self, session_id: str | None = None) -> KeyValueCache:
-        """
-        Retrieves the volatile key-value cache associated with the provided session.
-        :param session_id: The session to retrieve the volatile cache for. If not provided, the current context is used to find the session
-        :return: The volatile key-value cache.
-        """
-        if session_id is None:
-            session_id = Session.get_current_session_id()
-
-        if session_id is None or session_id == "":
-            raise Exception("No current session context available to retrieve volatile cache.")
-        return self._sessions.load(session_id).get_volatile_cache()
-
-    def get_non_volatile_cache(self, session_id: str | None = None) -> KeyValueCache:
-        """
-        Retrieves the non-volatile key-value cache associated with the provided session.
-        :param session_id: The session to retrieve the non-volatile cache for. If not provided, the current context is used to find the session
-        :return: The non-volatile key-value cache.
-        """
-        if session_id is None:
-            session_id = Session.get_current_session_id()
-
-        if session_id is None or session_id == "":
-            raise Exception("No current session context available to retrieve non-volatile cache.")
-        return self._sessions.load(session_id).get_non_volatile_cache()
-
 
 class GlobalRuntime(Runtime, metaclass=Singleton):
     """
@@ -204,66 +201,37 @@ def instance() -> Runtime:
         return GlobalRuntime()
 
 
-class ModuleLoader:
+class AuxiliaryCache:
     """
-    ModuleLoader is responsible for loading agent modules dynamically.
+    AuxiliaryCache provides access to volatile and non-volatile key-value caches associated with
+    the current or a provided session.
     """
 
-    _runtime: Optional[Runtime] = None
-    _lock: RLock = RLock()
-
     @staticmethod
-    def runtime() -> Runtime:
+    def get_volatile_cache(session_id: str | None = None) -> KeyValueCache:
         """
-        Return the Runtime instance set to load the module. By default this is the
-        global singleton Runtime instance.
-        """
-        return ModuleLoader._runtime or GlobalRuntime.instance()
-
-    @staticmethod
-    def attach(runtime: Runtime):
+        Retrieves the volatile key-value cache associated with the provided session.
+        :param session_id: The session to retrieve the volatile cache for. If not provided, the current context is used to find the session
+        :return: The volatile key-value cache.
         """
-        Attach a Runtime instance to the ModuleLoader.
+        if session_id is None:
+            session_id = Session.get_current_session_id()
 
-        This method sets the Runtime instance that will be used by the ModuleLoader for
-        loading and managing modules. It ensures thread-safety using a lock and validates
-        that only one Runtime can be attached at a time.
+        if session_id is None or session_id == "":
+            raise Exception("No current session context available to retrieve volatile cache.")
 
-        :param runtime: The Runtime instance to attach to the ModuleLoader.
-        :raises Exception: If a different runtime instance is already attached to the ModuleLoader.
-        """
-        with ModuleLoader._lock:
-            if ModuleLoader._runtime is not None and ModuleLoader._runtime != runtime:
-                raise Exception("A different runtime is already attached")
-            ModuleLoader._runtime = runtime
+        return Runtime.current().sessions().load(session_id).get_volatile_cache()
 
     @staticmethod
-    def detach(runtime: Runtime):
+    def get_non_volatile_cache(session_id: str | None = None) -> KeyValueCache:
         """
-        Detach a Runtime instance from the ModuleLoader.
-
-        This method removes the Runtime association from the ModuleLoader in a thread-safe manner.
-        It validates that the runtime being detached matches the currently attached runtime before
-        proceeding with the detachment.
-
-        :param runtime: The runtime instance to detach from the ModuleLoader.
-        :raises Exception: If a different runtime is currently attached than the one being detached.
-        """
-        with ModuleLoader._lock:
-            if ModuleLoader._runtime is not None and ModuleLoader._runtime != runtime:
-                raise Exception("A different runtime is already attached")
-            ModuleLoader._runtime = None
-
-    @staticmethod
-    def load(runtime: Runtime, module: str) -> ModuleType:
+        Retrieves the non-volatile key-value cache associated with the provided session.
+        :param session_id: The session to retrieve the non-volatile cache for. If not provided, the current context is used to find the session
+        :return: The non-volatile key-value cache.
         """
-        Load a module within the context of a given runtime.
-        :param runtime: The runtime environment to be associated with the module loading process.
-        :param module: The name of the module to import (e.g., 'os.path' or 'mypackage.mymodule').
-        :return: The imported module object.
+        if session_id is None:
+            session_id = Session.get_current_session_id()
 
-        :raises ModuleNotFoundError: If the specified module cannot be found.
-        :raises ImportError: If there's an error during the module import process.
-        """
-        with ModuleLoader._lock, runtime:
-            return importlib.import_module(module)
+        if session_id is None or session_id == "":
+            raise Exception("No current session context available to retrieve non-volatile cache.")
+        return Runtime.current().sessions().load(session_id).get_non_volatile_cache()

ak-py/src/agentkernel/core/service.py

@@ -1,10 +1,8 @@
 import logging
 import uuid
-from typing import Any
 
-from agentkernel.core.model import AgentReply, AgentReplyText, AgentRequestText
-
-from ..core import Agent, AgentRequest, GlobalRuntime, Runtime, Session
+from ..core import Agent, AgentRequest, Runtime, Session
+from ..core.model import AgentReply, AgentReplyText, AgentRequestText
 
 
 class AgentService:
@@ -17,7 +15,7 @@ def __init__(self):
         self._log = logging.getLogger("ak.core.service.agentservice")
         self._agent = None
         self._session = None
-        self._runtime = GlobalRuntime.instance()
+        self._runtime = Runtime.current()
 
     @property
     def runtime(self) -> Runtime:

ak-py/src/agentkernel/core/session/base.py

@@ -86,7 +86,7 @@ def set(self, session: Session) -> None:
         is at capacity, the least recently used session is removed before adding
         the new session. In either case the session is marked as most recently used.
 
-        :param session (Session): The session object to store in the cache.
+        :param session: The session object to be stored in the cache.
         """
         with self._lock:
             if session.id in self._cache: