refactor: improve caching system and optimize extension/plugin path resolution
- Change cache key type from `str` to `Any` in cache helper functions - Add `determine_cache_key` helper to generate consistent cache keys from agent profile and project - Add extension call logging with `EXTENSIONS_LOG` environment variable support - Implement caching for extension classes, enabled plugins, and subagent paths - Optimize `get_abs_path` with `_resolve_path` helper to avoid redundant `get_base_dir` calls - Store
frdel committed
Mar 18, 2026 at 12:40 UTC
d6f67b4df31169acd514abc6968f9f349e33e09f
8 files changed
+154
-21
agent.py
+1
-1
@@ -176,7 +176,7 @@ class AgentContext:
176
# recursive is not used now, prepared for context hierarchy
177
self.output_data[key] = value
178
179
- @extension.extensible
179
+ # @extension.extensible
180
def output(self):
181
return {
182
"id": self.id,
helpers/api.py
+1
-1
@@ -16,7 +16,7 @@ from helpers import files, cache
16
17
ThreadLockType = Union[threading.Lock, threading.RLock]
18
19
-CACHE_AREA = "api_handlers(api)(plugins)(extensions)"
19
+CACHE_AREA = "api_handlers(api)"
20
cache.toggle_area(CACHE_AREA, False) # cache off for now
21
22
Input = dict
helpers/cache.py
+12
-4
@@ -18,14 +18,14 @@ def toggle_area(area: str, enabled: bool) -> None:
18
_enabled_areas[area] = enabled
19
20
21
-def has(area: str, key: str) -> bool:
21
+def has(area: str, key: Any) -> bool:
22
if not _is_enabled(area):
23
return False
24
with _lock:
25
return key in _cache.get(area, {})
26
27
28
-def add(area: str, key: str, data: Any) -> None:
28
+def add(area: str, key: Any, data: Any) -> None:
29
if not _is_enabled(area):
30
return
31
with _lock:
@@ -34,14 +34,14 @@ def add(area: str, key: str, data: Any) -> None:
34
_cache[area][key] = data
35
36
37
-def get(area: str, key: str, default: Any = None) -> Any:
37
+def get(area: str, key: Any, default: Any = None) -> Any:
38
if not _is_enabled(area):
39
return default
40
with _lock:
41
return _cache.get(area, {}).get(key, default)
42
43
44
-def remove(area: str, key: str) -> None:
44
+def remove(area: str, key: Any) -> None:
45
if not _is_enabled(area):
46
return
47
with _lock:
@@ -69,3 +69,11 @@ def _is_enabled(area: str) -> bool:
69
if not _enabled_global:
70
return False
71
return _enabled_areas.get(area, True)
72
+
73
+
74
+def determine_cache_key(agent, *additional):
75
+ if agent:
76
+ profile = agent.config.profile or "none"
77
+ project = agent.context.get_data("project") or "none"
78
+ return (profile, project, *additional)
79
+ return ("none", "none", *additional)
\ No newline at end of file
helpers/extension.py
+36
-4
@@ -5,6 +5,7 @@ from helpers import cache, subagents
5
from typing import TYPE_CHECKING
6
from functools import wraps
7
import inspect
8
+import os
9
10
if TYPE_CHECKING:
11
from agent import Agent
@@ -13,8 +14,10 @@ if TYPE_CHECKING:
14
DEFAULT_EXTENSIONS_FOLDER = "python/extensions"
15
USER_EXTENSIONS_FOLDER = "usr/extensions"
16
16
-_CACHE_AREA = "extension_folder_classes(extensions)(plugins)"
17
-cache.toggle_area(_CACHE_AREA, False) # cache off for now
17
+_EXTENSIONS_CACHE_AREA = "extension_folder_classes(extensions)"
18
+_CLASSES_CACHE_AREA = "extension_classes(extensions)"
19
+cache.toggle_area(_EXTENSIONS_CACHE_AREA, False)
20
+# cache.toggle_area(_CLASSES_CACHE_AREA, False)
21
22
23
class _Unset:
@@ -22,6 +25,24 @@ class _Unset:
25
26
27
_UNSET = _Unset()
28
+_EXTENSIONS_LOG_COUNTS: dict[str, int] = {}
29
+
30
+# debug - extensions call counter
31
+def _log_extension_call(name: str):
32
+ try:
33
+ every = int(os.getenv("EXTENSIONS_LOG", "0"))
34
+ except ValueError:
35
+ return
36
+
37
+ if every <= 0:
38
+ return
39
+
40
+ _EXTENSIONS_LOG_COUNTS[name] = _EXTENSIONS_LOG_COUNTS.get(name, 0) + 1
41
+ _EXTENSIONS_LOG_COUNTS["_total"] = _EXTENSIONS_LOG_COUNTS.get("_total", 0) + 1
42
+
43
+ if _EXTENSIONS_LOG_COUNTS["_total"] % every == 0:
44
+ for key, count in _EXTENSIONS_LOG_COUNTS.items():
45
+ print(f"{str(count):<6} {key}")
46
47
48
# decorator to enable implicit extension points in existing functions
@@ -80,6 +101,7 @@ def extensible(func):
101
102
start_point = f"{module_name}_{qual_name}_start"
103
end_point = f"{module_name}_{qual_name}_end"
104
+
105
agent = _get_agent(args, kwargs)
106
107
data = {
@@ -177,6 +199,8 @@ class Extension:
199
async def call_extensions_async(
200
extension_point: str, agent: "Agent|None" = None, **kwargs
201
):
202
+ _log_extension_call(extension_point)
203
+
204
# fetch classes for this extension point and agent
205
classes = _get_extension_classes(extension_point, agent=agent, **kwargs)
206
@@ -188,6 +212,8 @@ async def call_extensions_async(
212
213
214
def call_extensions_sync(extension_point: str, agent: "Agent|None" = None, **kwargs):
215
+ _log_extension_call(extension_point)
216
+
217
# fetch classes for this extension point and agent
218
classes = _get_extension_classes(extension_point, agent=agent, **kwargs)
219
@@ -230,6 +256,11 @@ def get_webui_extensions(
256
def _get_extension_classes(
257
extension_point: str, agent: "Agent|None" = None, **kwargs
258
) -> list[Type[Extension]]:
259
+ cache_key = cache.determine_cache_key(agent, extension_point)
260
+ cached = cache.get(_CLASSES_CACHE_AREA, cache_key)
261
+ if cached is not None:
262
+ return cached
263
+
264
# search for extension folders in all agent's paths
265
paths = subagents.get_paths(agent, "extensions/python", extension_point)
266
@@ -244,6 +275,7 @@ def _get_extension_classes(
275
classes = sorted(
276
unique.values(), key=lambda cls: _get_file_from_module(cls.__module__)
277
)
278
+ cache.add(_CLASSES_CACHE_AREA, cache_key, classes)
279
return classes
280
281
@@ -253,7 +285,7 @@ def _get_file_from_module(module_name: str) -> str:
285
286
def _get_extensions(folder: str):
287
folder = files.get_abs_path(folder)
256
- cached = cache.get(_CACHE_AREA, folder)
288
+ cached = cache.get(_EXTENSIONS_CACHE_AREA, folder)
289
if cached is not None:
290
return cached
291
@@ -261,5 +293,5 @@ def _get_extensions(folder: str):
293
return []
294
295
classes = extract_tools.load_classes_from_folder(folder, "*", Extension)
264
- cache.add(_CACHE_AREA, folder, classes)
296
+ cache.add(_EXTENSIONS_CACHE_AREA, folder, classes)
297
return classes
\ No newline at end of file
helpers/files.py
+12
-9
@@ -1,7 +1,6 @@
1
from abc import ABC, abstractmethod
2
from fnmatch import fnmatch
3
import json
4
-from ntpath import isabs
4
import os
5
import re
6
import base64
@@ -19,7 +18,7 @@ PLUGINS_DIR = "plugins"
18
PROJECTS_DIR = "projects"
19
USER_DIR = "usr"
20
TEMP_DIR = "tmp"
22
-
21
+_base_dir = os.path.dirname(os.path.abspath(os.path.join(__file__, "../")))
22
23
class VariablesPlugin(ABC):
24
@abstractmethod
@@ -529,9 +528,15 @@ def make_dirs(relative_path: str):
528
os.makedirs(os.path.dirname(abs_path), exist_ok=True)
529
530
531
+def _resolve_path(*relative_paths):
532
+ if len(relative_paths) == 1 and os.path.isabs(relative_paths[0]):
533
+ return relative_paths[0]
534
+ return os.path.join(_base_dir, *relative_paths)
535
+
536
+
537
def get_abs_path(*relative_paths):
538
"Convert relative paths to absolute paths based on the base directory."
534
- return os.path.join(get_base_dir(), *relative_paths)
539
+ return _resolve_path(*relative_paths)
540
541
542
def get_abs_path_dockerized(*relative_paths):
@@ -574,24 +579,22 @@ def normalize_a0_path(path: str):
579
580
581
def exists(*relative_paths):
577
- path = get_abs_path(*relative_paths)
582
+ path = _resolve_path(*relative_paths)
583
return os.path.exists(path)
584
585
586
def is_file(*relative_paths):
582
- path = get_abs_path(*relative_paths)
587
+ path = _resolve_path(*relative_paths)
588
return os.path.isfile(path)
589
590
591
def is_dir(*relative_paths):
587
- path = get_abs_path(*relative_paths)
592
+ path = _resolve_path(*relative_paths)
593
return os.path.isdir(path)
594
595
596
def get_base_dir():
592
- # Get the base directory from the current file path
593
- base_dir = os.path.dirname(os.path.abspath(os.path.join(__file__, "../")))
594
- return base_dir
597
+ return _base_dir
598
599
600
def basename(path: str, suffix: str | None = None):
helpers/plugins.py
+22
@@ -57,6 +57,10 @@ TOGGLE_FILE_PATTERN = ".toggle-[01]"
57
58
HOOKS_SCRIPT = "hooks.py"
59
HOOKS_CACHE_AREA = "plugin_hooks(plugins)"
60
+PLUGINS_LIST_CACHE_AREA = "plugins_list(plugins)"
61
+ENABLED_PLUGINS_LIST_CACHE_AREA = "enabled_plugins(plugins)"
62
+ENABLED_PLUGINS_PATHS_CACHE_AREA = "enabled_plugins_paths(plugins)"
63
+
64
65
_last_frontend_reload_notification_at = 0.0
66
@@ -114,6 +118,8 @@ def after_plugin_change(plugin_names: list[str] | None = None):
118
119
def clear_plugin_cache():
120
cache.clear("*(plugins)*")
121
+ cache.clear("*(extensions)*")
122
+ cache.clear("*(api)*")
123
124
125
def get_plugin_roots(plugin_name: str = "") -> List[str]:
@@ -125,6 +131,9 @@ def get_plugin_roots(plugin_name: str = "") -> List[str]:
131
132
133
def get_plugins_list():
134
+ if cached := cache.get(PLUGINS_LIST_CACHE_AREA, ""):
135
+ return cached
136
+
137
result: list[str] = []
138
seen_names: set[str] = set()
139
for root in get_plugin_roots():
@@ -137,6 +146,8 @@ def get_plugins_list():
146
seen_names.add(dir.name)
147
result.append(dir.name)
148
result.sort(key=lambda p: Path(p).name)
149
+
150
+ cache.add(PLUGINS_LIST_CACHE_AREA, "", result)
151
return result
152
153
@@ -291,6 +302,9 @@ def get_plugin_paths(*subpaths: str) -> List[str]:
302
303
304
def get_enabled_plugin_paths(agent: Agent | None, *subpaths: str) -> List[str]:
305
+ if cached := cache.get(ENABLED_PLUGINS_PATHS_CACHE_AREA, cache.determine_cache_key(agent, *subpaths)):
306
+ return cached
307
+
308
enabled = get_enabled_plugins(agent)
309
paths: list[str] = []
310
@@ -307,10 +321,16 @@ def get_enabled_plugin_paths(agent: Agent | None, *subpaths: str) -> List[str]:
321
path_pattern = files.get_abs_path(base_dir, *subpaths)
322
paths.extend(files.find_existing_paths_by_pattern(path_pattern))
323
324
+
325
+ cache.add(ENABLED_PLUGINS_PATHS_CACHE_AREA, cache.determine_cache_key(agent, *subpaths), paths)
326
+
327
return paths
328
329
330
def get_enabled_plugins(agent: Agent | None):
331
+ if cached := cache.get(ENABLED_PLUGINS_LIST_CACHE_AREA, cache.determine_cache_key(agent)):
332
+ return cached
333
+
334
plugins = get_plugins_list()
335
active = []
336
@@ -344,6 +364,8 @@ def get_enabled_plugins(agent: Agent | None):
364
if enabled:
365
active.append(plugin)
366
367
+ cache.add(ENABLED_PLUGINS_LIST_CACHE_AREA, cache.determine_cache_key(agent), active)
368
+
369
return active
370
371
helpers/subagents.py
+20
-2
@@ -1,15 +1,18 @@
1
from helpers import files
2
+from helpers import cache
3
from helpers import yaml as yaml_helper
3
-from typing import TypedDict, TYPE_CHECKING
4
+from typing import TypedDict, TYPE_CHECKING, Literal
5
from pydantic import BaseModel, model_validator
6
import json
6
-from typing import Literal
7
import os
8
9
GLOBAL_DIR = "."
10
USER_DIR = "usr"
11
DEFAULT_AGENTS_DIR = "agents"
12
USER_AGENTS_DIR = "usr/agents"
13
+PATHS_CACHE_AREA = "subagent_paths(plugins)"
14
+
15
+cache.toggle_area(PATHS_CACHE_AREA, False)
16
17
type Origin = Literal["default", "user", "project", "plugin"]
18
@@ -345,6 +348,20 @@ def get_paths(
348
) -> list[str]:
349
"""Returns list of file paths for the given agent and subpaths, searched in order of priority:
350
project/agents/, project/, usr/agents/, plugin agents/, agents/, usr/, plugins/, default."""
351
+ cache_key = cache.determine_cache_key(
352
+ agent,
353
+ *subpaths,
354
+ must_exist_completely,
355
+ include_project,
356
+ include_user,
357
+ include_default,
358
+ include_plugins,
359
+ default_root,
360
+ )
361
+ cached = cache.get(PATHS_CACHE_AREA, cache_key)
362
+ if cached is not None:
363
+ return cached
364
+
365
paths: list[str] = []
366
check_subpaths = subpaths if must_exist_completely else []
367
profile_name = agent.config.profile if agent and agent.config.profile else ""
@@ -411,6 +428,7 @@ def get_paths(
428
if (not must_exist_completely) or files.exists(path):
429
paths.append(path)
430
431
+ cache.add(PATHS_CACHE_AREA, cache_key, paths)
432
return paths
433
434
tests/test_extensions_stress.py
new
+50
@@ -0,0 +1,50 @@
1
+import cProfile
2
+import io
3
+import pstats
4
+import sys
5
+from pathlib import Path
6
+
7
+import pytest
8
+
9
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
10
+if str(PROJECT_ROOT) not in sys.path:
11
+ sys.path.insert(0, str(PROJECT_ROOT))
12
+
13
+from agent import Agent, AgentContext
14
+from helpers.extension import extensible
15
+from initialize import initialize_agent
16
+
17
+
18
+class PerfAgent(Agent):
19
+ @extensible
20
+ def perf_hook(self, value: int):
21
+ return value + 1
22
+
23
+
24
+@pytest.mark.parametrize("iterations", [10000])
25
+def test_extensible_method_performance_trace(iterations: int):
26
+ agent = PerfAgent(number=0, config=initialize_agent())
27
+ context = agent.context
28
+
29
+ try:
30
+ profiler = cProfile.Profile()
31
+ profiler.enable()
32
+
33
+ result = 0
34
+ for i in range(iterations):
35
+ result = agent.perf_hook(i)
36
+
37
+ profiler.disable()
38
+
39
+ output = io.StringIO()
40
+ stats = pstats.Stats(profiler, stream=output)
41
+ stats.sort_stats("cumulative")
42
+ stats.print_stats(20)
43
+
44
+ print(f"\n[extensible perf] iterations={iterations} result={result}")
45
+ print(output.getvalue())
46
+
47
+ assert result == iterations
48
+ finally:
49
+ if context:
50
+ AgentContext.remove(context.id)