Add extensible decorator and cache controls
Introduce an extensible extension point system and finer cache control. Added extensible decorator in python/helpers/extension.py to wrap sync/async functions, emit start/end extension points, detect Agent instances, and allow extensions to modify inputs/results/exceptions. Reworked extension loading to use the cache API and a dedicated cache area. Enhanced python/helpers/cache.py with global and per-area toggles and early-return behavior in add/get/remove. Updated API handler caching (python/helpers/api.py) to use the cache area toggle and simplify cache logic. Marked several initialization functions as @extension.extensible in initialize.py. Moved nest_asyncio.apply() from agent.py into python/helpers/runtime.py to ensure the event loop patch is applied at runtime initialization. Also applied small typing/formatting tweaks and path handling improvements.
frdel committed
Feb 25, 2026 at 16:52 UTC
d82b121bc9a124dae796e49502aca77efc893536
6 files changed
+169
-24
agent.py
-3
@@ -1,7 +1,4 @@
1
import asyncio, random, string, threading
2
-import nest_asyncio
3
-
4
-nest_asyncio.apply()
2
3
from collections import OrderedDict
4
from dataclasses import dataclass, field
initialize.py
+7
-1
@@ -1,9 +1,10 @@
1
from agent import AgentConfig
2
import models
3
-from python.helpers import runtime, settings, defer
3
+from python.helpers import runtime, settings, defer, extension
4
from python.helpers.print_style import PrintStyle
5
6
7
+@extension.extensible
8
def initialize_agent(override_settings: dict | None = None):
9
current_settings = settings.get_settings()
10
if override_settings:
@@ -118,12 +119,14 @@ def initialize_agent(override_settings: dict | None = None):
119
# return config object
120
return config
121
122
+@extension.extensible
123
def initialize_chats():
124
from python.helpers import persist_chat
125
async def initialize_chats_async():
126
persist_chat.load_tmp_chats()
127
return defer.DeferredTask().start_task(initialize_chats_async)
128
129
+@extension.extensible
130
def initialize_mcp():
131
set = settings.get_settings()
132
async def initialize_mcp_async():
@@ -131,14 +134,17 @@ def initialize_mcp():
134
return _initialize_mcp(set["mcp_servers"])
135
return defer.DeferredTask().start_task(initialize_mcp_async)
136
137
+@extension.extensible
138
def initialize_job_loop():
139
from python.helpers.job_loop import run_loop
140
return defer.DeferredTask("JobLoop").start_task(run_loop)
141
142
+@extension.extensible
143
def initialize_preload():
144
import preload
145
return defer.DeferredTask().start_task(preload.preload)
146
147
+@extension.extensible
148
def initialize_migration():
149
from python.helpers import migration, dotenv
150
# run migration
python/helpers/api.py
+5
-8
@@ -17,8 +17,7 @@ from python.helpers import files, cache
17
ThreadLockType = Union[threading.Lock, threading.RLock]
18
19
CACHE_AREA = "api_handlers(api)(plugins)"
20
-CACHE_ENABLED = False
21
-
20
+cache.toggle_area(CACHE_AREA, False) # cache off for now
21
22
Input = dict
23
Output = Union[Dict[str, Any], Response, TypedDict] # type: ignore
@@ -210,10 +209,9 @@ def register_api_route(app: Flask, lock: ThreadLockType) -> None:
209
210
async def _dispatch(path: str) -> BaseResponse:
211
# Return cached wrapped handler if available
213
- if CACHE_ENABLED:
214
- cached = cache.get(CACHE_AREA, path)
215
- if cached is not None:
216
- return await cached()
212
+ cached = cache.get(CACHE_AREA, path)
213
+ if cached is not None:
214
+ return await cached()
215
216
# Resolve file path for the handler
217
# Try built-in api folder first, then plugin api folders
@@ -261,8 +259,7 @@ def register_api_route(app: Flask, lock: ThreadLockType) -> None:
259
if handler_cls.requires_loopback():
260
handler_fn = requires_loopback(handler_fn)
261
264
- if CACHE_ENABLED:
265
- cache.add(CACHE_AREA, path, handler_fn)
262
+ cache.add(CACHE_AREA, path, handler_fn)
263
return await handler_fn()
264
265
app.add_url_rule(
python/helpers/cache.py
+24
@@ -5,8 +5,22 @@ from typing import Any
5
_lock = threading.RLock()
6
_cache: dict[str, dict[str, Any]] = {}
7
8
+_enabled_global: bool = True
9
+_enabled_areas: dict[str, bool] = {}
10
+
11
+
12
+def toggle_global(enabled: bool) -> None:
13
+ global _enabled_global
14
+ _enabled_global = enabled
15
+
16
+
17
+def toggle_area(area: str, enabled: bool) -> None:
18
+ _enabled_areas[area] = enabled
19
+
20
21
def add(area: str, key: str, data: Any) -> None:
22
+ if not _is_enabled(area):
23
+ return
24
with _lock:
25
if area not in _cache:
26
_cache[area] = {}
@@ -14,11 +28,15 @@ def add(area: str, key: str, data: Any) -> None:
28
29
30
def get(area: str, key: str, default: Any = None) -> Any:
31
+ if not _is_enabled(area):
32
+ return default
33
with _lock:
34
return _cache.get(area, {}).get(key, default)
35
36
37
def remove(area: str, key: str) -> None:
38
+ if not _is_enabled(area):
39
+ return
40
with _lock:
41
if area in _cache:
42
_cache[area].pop(key, None)
@@ -38,3 +56,9 @@ def clear(area: str) -> None:
56
def clear_all() -> None:
57
with _lock:
58
_cache.clear()
59
+
60
+
61
+def _is_enabled(area: str) -> bool:
62
+ if not _enabled_global:
63
+ return False
64
+ return _enabled_areas.get(area, True)
python/helpers/extension.py
+130
-12
@@ -1,7 +1,11 @@
1
from abc import abstractmethod
2
from typing import Any
3
from python.helpers import extract_tools, files
4
+from python.helpers import cache
5
from typing import TYPE_CHECKING
6
+from functools import wraps
7
+import asyncio
8
+import inspect
9
10
if TYPE_CHECKING:
11
from agent import Agent
@@ -10,13 +14,123 @@ if TYPE_CHECKING:
14
DEFAULT_EXTENSIONS_FOLDER = "python/extensions"
15
USER_EXTENSIONS_FOLDER = "usr/extensions"
16
13
-_cache: dict[str, list[type["Extension"]]] = {}
17
+_CACHE_AREA = "extension_folder_classes(extensions)(plugins)"
18
+cache.toggle_area(_CACHE_AREA, False) # cache off for now
19
+
20
+
21
+class _Unset:
22
+ pass
23
+
24
+
25
+_UNSET = _Unset()
26
+
27
+
28
+# decorator to enable implicit extension points in existing functions
29
+def extensible(func):
30
+ """Make a function emit two implicit extension points around its execution.
31
+
32
+ The decorator derives two extension point names from the wrapped function:
33
+
34
+ - ``{func.__module__}.{func.__name__}-start``
35
+ - ``{func.__module__}.{func.__name__}-end``
36
+
37
+ When the wrapped function is called, the decorator builds a mutable ``data``
38
+ payload and passes it to both extension points via ``call_extensions``:
39
+
40
+ - ``data["args"]``: the original positional arguments tuple
41
+ - ``data["kwargs"]``: the original keyword arguments dict
42
+ - ``data["result"]``: initialized to an internal sentinel; extensions may
43
+ set this to short-circuit the wrapped function
44
+ - ``data["exception"]``: initialized to an internal sentinel; extensions may
45
+ set this to a ``BaseException`` instance to force-raise
46
+
47
+ Behavior:
48
+
49
+ - ``-start`` extensions run first and may mutate ``data["args"]`` /
50
+ ``data["kwargs"]``, set ``data["result"]`` to skip calling ``func``, or set
51
+ ``data["exception"]`` to abort by raising.
52
+ - If ``data["result"]`` is still unset, the decorator calls ``func`` (awaiting
53
+ it if it is async) and stores either the return value into ``data["result"]``
54
+ or the raised error into ``data["exception"]``.
55
+ - ``-end`` extensions run last and may further transform the outcome by
56
+ rewriting ``data["result"]`` or replacing/clearing ``data["exception"]``.
57
+
58
+ Finally, if ``data["exception"]`` contains an exception it is raised;
59
+ otherwise ``data["result"]`` is returned.
60
+ """
61
+ @wraps(func)
62
+ async def _inner_async(*args, **kwargs):
63
+ from agent import Agent
64
+
65
+ # prepare extension points data
66
+ module_name = getattr(func, "__module__", "")
67
+ func_name = getattr(func, "__name__", "")
68
+ start_point = f"{module_name}.{func_name}-start"
69
+ end_point = f"{module_name}.{func_name}-end"
70
+
71
+ # try to find agent instance for better extension determination
72
+ agent = kwargs.get("agent")
73
+ if (not agent or not isinstance(agent, Agent)) and args:
74
+ try:
75
+ for a in args:
76
+ if isinstance(a, Agent):
77
+ agent = a
78
+ break
79
+ except Exception:
80
+ agent = None
81
+
82
+ # build extension data object - func input/output
83
+ data = {
84
+ "args": args,
85
+ "kwargs": kwargs,
86
+ "result": _UNSET,
87
+ "exception": _UNSET,
88
+ }
89
+
90
+ # call start extensions, these can modify inputs, produce output or exception
91
+ await call_extensions(start_point, agent=agent, data=data)
92
+
93
+ # if there is an explicit exception set, raise it
94
+ exc = data.get("exception")
95
+ if isinstance(exc, BaseException):
96
+ raise exc
97
+
98
+ # if there is no result set, call the original function
99
+ if data.get("result") is _UNSET:
100
+ try:
101
+ if inspect.iscoroutinefunction(func):
102
+ data["result"] = await func(*args, **kwargs)
103
+ else:
104
+ data["result"] = func(*args, **kwargs)
105
+ except Exception as e:
106
+ data["exception"] = e
107
+
108
+ # call end extensions, these can modify outputs or exception
109
+ await call_extensions(end_point, agent=agent, data=data)
110
+
111
+ # if there's an exception, raise it
112
+ exc = data.get("exception")
113
+ if isinstance(exc, BaseException):
114
+ raise exc
115
+
116
+ # if there's a result, return it
117
+ result = data.get("result")
118
+ return None if result is _UNSET else result
119
+
120
+ if inspect.iscoroutinefunction(func):
121
+ return _inner_async
122
+
123
+ @wraps(func)
124
+ def _inner_sync(*args, **kwargs):
125
+ return asyncio.run(_inner_async(*args, **kwargs))
126
+
127
+ return _inner_sync
128
129
130
class Extension:
131
132
def __init__(self, agent: "Agent|None", **kwargs):
19
- self.agent: "Agent" = agent # type: ignore < here we ignore the type check as there are currently no extensions without an agent
133
+ self.agent: "Agent|None" = agent
134
self.kwargs = kwargs
135
136
@abstractmethod
@@ -30,10 +144,14 @@ async def call_extensions(
144
from python.helpers import projects, subagents, plugins
145
146
# search for extension folders in all agent's paths
33
- paths = subagents.get_paths(agent, "extensions", extension_point, default_root="python")
147
+ paths = subagents.get_paths(
148
+ agent, "extensions", extension_point, default_root="python"
149
+ )
150
151
# Add plugin backend extension paths (plugins/*/extensions/python/{extension_point})
36
- plugin_paths = plugins.get_enabled_plugin_paths(agent, "extensions", "python", extension_point)
152
+ plugin_paths = plugins.get_enabled_plugin_paths(
153
+ agent, "extensions", "python", extension_point
154
+ )
155
paths.extend(p for p in plugin_paths if p not in paths)
156
157
all_exts = [cls for path in paths for cls in _get_extensions(path)]
@@ -58,14 +176,14 @@ def _get_file_from_module(module_name: str) -> str:
176
177
178
def _get_extensions(folder: str):
61
- global _cache
179
folder = files.get_abs_path(folder)
63
- if folder in _cache:
64
- classes = _cache[folder]
65
- else:
66
- if not files.exists(folder):
67
- return []
68
- classes = extract_tools.load_classes_from_folder(folder, "*", Extension)
69
- _cache[folder] = classes
180
+ cached = cache.get(_CACHE_AREA, folder)
181
+ if cached is not None:
182
+ return cached
183
+
184
+ if not files.exists(folder):
185
+ return []
186
187
+ classes = extract_tools.load_classes_from_folder(folder, "*", Extension)
188
+ cache.add(_CACHE_AREA, folder, classes)
189
return classes
python/helpers/runtime.py
+3
@@ -8,6 +8,9 @@ import asyncio
8
import threading
9
import queue
10
import sys
11
+import nest_asyncio
12
+
13
+nest_asyncio.apply()
14
15
T = TypeVar("T")
16
R = TypeVar("R")