| 1 | from abc import abstractmethod |
| 2 | from typing import Any, Awaitable, Type, cast |
| 3 | from helpers import modules, files |
| 4 | from helpers import cache |
| 5 | from typing import TYPE_CHECKING |
| 6 | from functools import wraps |
| 7 | import inspect |
| 8 | import os |
| 9 | |
| 10 | from helpers.print_style import PrintStyle |
| 11 | |
| 12 | if TYPE_CHECKING: |
| 13 | from agent import Agent |
| 14 | |
| 15 | |
| 16 | DEFAULT_EXTENSIONS_FOLDER = "python/extensions" |
| 17 | USER_EXTENSIONS_FOLDER = "usr/extensions" |
| 18 | |
| 19 | _EXTENSIONS_CACHE_AREA = "extension_folder_classes(extensions)" |
| 20 | _CLASSES_CACHE_AREA = "extension_classes(extensions)" |
| 21 | _WEBUI_MANIFEST_CACHE_AREA = "webui_extension_manifest(extensions)(plugins)" |
| 22 | _WEBUI_MANIFEST_SUFFIXES = { |
| 23 | "html": (".html", ".htm", ".xhtml"), |
| 24 | "js": (".js", ".mjs"), |
| 25 | } |
| 26 | # cache.toggle_area(_EXTENSIONS_CACHE_AREA, False) |
| 27 | # cache.toggle_area(_CLASSES_CACHE_AREA, False) |
| 28 | |
| 29 | |
| 30 | class _Unset: |
| 31 | pass |
| 32 | |
| 33 | |
| 34 | _UNSET = _Unset() |
| 35 | _EXTENSIONS_LOG_COUNTS: dict[str, int] = {} |
| 36 | |
| 37 | |
| 38 | # debug - extensions call counter |
| 39 | def _log_extension_call(name: str): |
| 40 | try: |
| 41 | every = int(os.getenv("EXTENSIONS_LOG", "0")) |
| 42 | except ValueError: |
| 43 | return |
| 44 | |
| 45 | if every <= 0: |
| 46 | return |
| 47 | |
| 48 | _EXTENSIONS_LOG_COUNTS[name] = _EXTENSIONS_LOG_COUNTS.get(name, 0) + 1 |
| 49 | _EXTENSIONS_LOG_COUNTS["_total"] = _EXTENSIONS_LOG_COUNTS.get("_total", 0) + 1 |
| 50 | |
| 51 | if _EXTENSIONS_LOG_COUNTS["_total"] % every == 0: |
| 52 | for key, count in _EXTENSIONS_LOG_COUNTS.items(): |
| 53 | print(f"{str(count):<6} {key}") |
| 54 | |
| 55 | |
| 56 | # decorator to enable implicit extension points in existing functions |
| 57 | def extensible(func): |
| 58 | """Make a function emit two implicit extension points around its execution. |
| 59 | |
| 60 | The decorator derives two extension point folder paths from the wrapped |
| 61 | function: |
| 62 | |
| 63 | - ``_functions/<module path>/<qualname path>/start`` |
| 64 | - ``_functions/<module path>/<qualname path>/end`` |
| 65 | |
| 66 | Module path segments come from ``func.__module__`` split by ``.``. |
| 67 | Qualname path segments come from the full nested ``func.__qualname__`` split |
| 68 | by ``.``, excluding ``<locals>``. |
| 69 | |
| 70 | Example: |
| 71 | |
| 72 | - module ``helpers.something`` |
| 73 | - qualname ``Outer.Inner.__init__`` |
| 74 | |
| 75 | becomes: |
| 76 | |
| 77 | - ``_functions/helpers/something/Outer/Inner/__init__/start`` |
| 78 | - ``_functions/helpers/something/Outer/Inner/__init__/end`` |
| 79 | |
| 80 | When the wrapped function is called, the decorator builds a mutable ``data`` |
| 81 | payload and passes it to both extension points: |
| 82 | |
| 83 | - ``data["args"]``: positional args (extensions may replace/mutate) |
| 84 | - ``data["kwargs"]``: keyword args (extensions may replace/mutate) |
| 85 | - ``data["result"]``: initialized to an internal sentinel; extensions may set |
| 86 | this to short-circuit the wrapped function |
| 87 | - ``data["exception"]``: initialized to an internal sentinel; extensions may |
| 88 | set this to a ``BaseException`` instance to force-raise |
| 89 | |
| 90 | Sync functions call ``call_extensions_sync``. Async functions call |
| 91 | ``call_extensions_async``. |
| 92 | |
| 93 | Behavior: |
| 94 | |
| 95 | - ``start`` extensions run first and may mutate inputs or set |
| 96 | ``data["result"]`` / ``data["exception"]``. |
| 97 | - If ``data["result"]`` is still unset, the decorator calls the wrapped |
| 98 | function using the possibly modified ``data["args"]`` / ``data["kwargs"]``. |
| 99 | - ``end`` extensions run last and may rewrite ``data["result"]`` or replace / |
| 100 | clear ``data["exception"]``. |
| 101 | |
| 102 | Finally, if ``data["exception"]`` contains an exception it is raised; |
| 103 | otherwise ``data["result"]`` is returned. |
| 104 | """ |
| 105 | |
| 106 | def _get_agent(args, kwargs): |
| 107 | from agent import Agent |
| 108 | |
| 109 | candidate = kwargs.get("agent") |
| 110 | if isinstance(candidate, Agent) and bool(getattr(candidate, "__dict__", None)): |
| 111 | return candidate |
| 112 | |
| 113 | for a in args: |
| 114 | if isinstance(a, Agent) and bool(getattr(a, "__dict__", None)): |
| 115 | return a |
| 116 | |
| 117 | return None |
| 118 | |
| 119 | def _prepare_inputs(args, kwargs): |
| 120 | module_name = getattr(func, "__module__", "") |
| 121 | qual_name = getattr(func, "__qualname__", "") |
| 122 | if not module_name or not qual_name: |
| 123 | return None |
| 124 | |
| 125 | module_parts = [part for part in module_name.split(".") if part] |
| 126 | qual_parts = [part for part in qual_name.split(".") if part and part != "<locals>"] |
| 127 | if not module_parts or not qual_parts: |
| 128 | return None |
| 129 | |
| 130 | base_path = os.path.join("_functions", *module_parts, *qual_parts) |
| 131 | start_point = os.path.join(base_path, "start") |
| 132 | end_point = os.path.join(base_path, "end") |
| 133 | |
| 134 | agent = _get_agent(args, kwargs) |
| 135 | |
| 136 | data = { |
| 137 | "args": args, |
| 138 | "kwargs": kwargs, |
| 139 | "result": _UNSET, |
| 140 | "exception": None, |
| 141 | } |
| 142 | |
| 143 | return start_point, end_point, agent, data |
| 144 | |
| 145 | def _process_result(data): |
| 146 | exc = data.get("exception") |
| 147 | if isinstance(exc, BaseException): |
| 148 | raise exc |
| 149 | |
| 150 | return data.get("result") |
| 151 | |
| 152 | def _call_original(data): |
| 153 | call_args = data.get("args") |
| 154 | call_kwargs = data.get("kwargs") |
| 155 | |
| 156 | if not isinstance(call_args, tuple): |
| 157 | call_args = (call_args,) |
| 158 | if not isinstance(call_kwargs, dict): |
| 159 | call_kwargs = {} |
| 160 | |
| 161 | try: |
| 162 | data["result"] = func(*call_args, **call_kwargs) |
| 163 | except Exception as e: |
| 164 | data["exception"] = e |
| 165 | return _UNSET |
| 166 | |
| 167 | async def _run_async(*args, **kwargs): |
| 168 | prepared = _prepare_inputs(args, kwargs) |
| 169 | if prepared is None: |
| 170 | return await func(*args, **kwargs) |
| 171 | |
| 172 | start_point, end_point, agent, data = prepared |
| 173 | |
| 174 | # call pre-extensions |
| 175 | await call_extensions_async(start_point, agent=agent, data=data) |
| 176 | |
| 177 | # call the original if pre-extensions don't return a result |
| 178 | if (result := _process_result(data)) is _UNSET: |
| 179 | _call_original(data) |
| 180 | try: |
| 181 | data["result"] = await data["result"] |
| 182 | except Exception as e: |
| 183 | data["exception"] = e |
| 184 | |
| 185 | # call post-extensions |
| 186 | await call_extensions_async(end_point, agent=agent, data=data) |
| 187 | |
| 188 | result = _process_result(data) |
| 189 | return None if result is _UNSET else result |
| 190 | |
| 191 | def _run_sync(*args, **kwargs): |
| 192 | prepared = _prepare_inputs(args, kwargs) |
| 193 | if prepared is None: |
| 194 | return func(*args, **kwargs) |
| 195 | |
| 196 | start_point, end_point, agent, data = prepared |
| 197 | |
| 198 | # call pre-extensions |
| 199 | call_extensions_sync(start_point, agent=agent, data=data) |
| 200 | |
| 201 | # call the original if pre-extensions don't return a result |
| 202 | if (result := _process_result(data)) is _UNSET: |
| 203 | _call_original(data) |
| 204 | |
| 205 | # call post-extensions |
| 206 | call_extensions_sync(end_point, agent=agent, data=data) |
| 207 | |
| 208 | result = _process_result(data) |
| 209 | return None if result is _UNSET else result |
| 210 | |
| 211 | if inspect.iscoroutinefunction(func): |
| 212 | return wraps(func)(_run_async) |
| 213 | |
| 214 | return wraps(func)(_run_sync) |
| 215 | |
| 216 | |
| 217 | class Extension: |
| 218 | |
| 219 | def __init__(self, agent: "Agent|None", **kwargs): |
| 220 | self.agent: "Agent|None" = agent |
| 221 | self.kwargs = kwargs |
| 222 | |
| 223 | @abstractmethod |
| 224 | def execute(self, **kwargs) -> None | Awaitable[None]: |
| 225 | pass |
| 226 | |
| 227 | |
| 228 | async def call_extensions_async( |
| 229 | extension_point: str, agent: "Agent|None" = None, **kwargs |
| 230 | ): |
| 231 | _log_extension_call(extension_point) |
| 232 | |
| 233 | # fetch classes for this extension point and agent |
| 234 | classes = _get_extension_classes(extension_point, agent=agent, **kwargs) |
| 235 | |
| 236 | # execute unique extensions |
| 237 | for cls in classes: |
| 238 | result = cls(agent=agent).execute(**kwargs) |
| 239 | if isinstance(result, Awaitable): |
| 240 | await result |
| 241 | |
| 242 | |
| 243 | def call_extensions_sync(extension_point: str, agent: "Agent|None" = None, **kwargs): |
| 244 | _log_extension_call(extension_point) |
| 245 | |
| 246 | # fetch classes for this extension point and agent |
| 247 | classes = _get_extension_classes(extension_point, agent=agent, **kwargs) |
| 248 | |
| 249 | # execute unique extensions |
| 250 | for cls in classes: |
| 251 | result = cls(agent=agent).execute(**kwargs) |
| 252 | if isinstance(result, Awaitable): |
| 253 | raise ValueError( |
| 254 | f"Extension {cls.__name__} returned awaitable in sync mode" |
| 255 | ) |
| 256 | |
| 257 | |
| 258 | def get_webui_extensions( |
| 259 | agent: "Agent | None", extension_point: str, filters: list[str] | None = None |
| 260 | ): |
| 261 | from helpers import subagents |
| 262 | |
| 263 | entries: list[str] = [] |
| 264 | effective_filters = filters or ["*"] |
| 265 | |
| 266 | # search for extension folders in all agent's paths |
| 267 | folders = subagents.get_paths( |
| 268 | agent, |
| 269 | "extensions/webui", |
| 270 | extension_point, |
| 271 | ) |
| 272 | |
| 273 | extensions = [] |
| 274 | |
| 275 | for folder in folders: |
| 276 | for filter in effective_filters: |
| 277 | pattern = files.get_abs_path(folder, filter) |
| 278 | extensions.extend(files.find_existing_paths_by_pattern(pattern)) |
| 279 | |
| 280 | for extension in extensions: |
| 281 | rel_path = files.deabsolute_path(extension) |
| 282 | entries.append(rel_path) |
| 283 | |
| 284 | return entries |
| 285 | |
| 286 | |
| 287 | def get_webui_extension_manifest( |
| 288 | agent: "Agent | None", |
| 289 | ) -> dict[str, dict[str, list[str]]]: |
| 290 | """Return every WebUI extension URL grouped by asset type and extension point.""" |
| 291 | from helpers import subagents |
| 292 | |
| 293 | cache_key = cache.determine_cache_key(agent) |
| 294 | cached = cache.get(_WEBUI_MANIFEST_CACHE_AREA, cache_key) |
| 295 | if cached is not None: |
| 296 | return cached |
| 297 | |
| 298 | manifest: dict[str, dict[str, list[str]]] = { |
| 299 | asset_type: {} for asset_type in _WEBUI_MANIFEST_SUFFIXES |
| 300 | } |
| 301 | roots = subagents.get_paths(agent, "extensions/webui") |
| 302 | for root in roots: |
| 303 | relative_files = sorted(files.list_files_in_dir_recursively(root)) |
| 304 | for asset_type, suffixes in _WEBUI_MANIFEST_SUFFIXES.items(): |
| 305 | for suffix in suffixes: |
| 306 | for relative_file in relative_files: |
| 307 | if not relative_file.lower().endswith(suffix): |
| 308 | continue |
| 309 | extension_point = os.path.dirname(relative_file).replace( |
| 310 | os.sep, "/" |
| 311 | ) |
| 312 | if not extension_point or extension_point == ".": |
| 313 | continue |
| 314 | absolute_path = files.get_abs_path(root, relative_file) |
| 315 | relative_path = files.deabsolute_path(absolute_path).replace( |
| 316 | os.sep, "/" |
| 317 | ) |
| 318 | manifest[asset_type].setdefault(extension_point, []).append( |
| 319 | "/" + relative_path.lstrip("/") |
| 320 | ) |
| 321 | |
| 322 | cache.add(_WEBUI_MANIFEST_CACHE_AREA, cache_key, manifest) |
| 323 | return manifest |
| 324 | |
| 325 | |
| 326 | def _get_extension_classes( |
| 327 | extension_point: str, agent: "Agent|None" = None, **kwargs |
| 328 | ) -> list[Type[Extension]]: |
| 329 | from helpers import subagents |
| 330 | |
| 331 | cache_key = cache.determine_cache_key(agent, extension_point) |
| 332 | cached = cache.get(_CLASSES_CACHE_AREA, cache_key) |
| 333 | if cached is not None: |
| 334 | return cached |
| 335 | |
| 336 | # search for extension folders in all agent's paths |
| 337 | paths = subagents.get_paths(agent, "extensions/python", extension_point) |
| 338 | |
| 339 | all_exts = [cls for path in paths for cls in _get_extensions(path)] |
| 340 | |
| 341 | # merge: first ocurrence of file name is the override |
| 342 | unique = {} |
| 343 | for cls in all_exts: |
| 344 | file = _get_file_from_module(cls.__module__) |
| 345 | if file not in unique: |
| 346 | unique[file] = cls |
| 347 | classes = sorted( |
| 348 | unique.values(), key=lambda cls: _get_file_from_module(cls.__module__) |
| 349 | ) |
| 350 | cache.add(_CLASSES_CACHE_AREA, cache_key, classes) |
| 351 | return classes |
| 352 | |
| 353 | |
| 354 | def _get_file_from_module(module_name: str) -> str: |
| 355 | return module_name.split(".")[-1] |
| 356 | |
| 357 | |
| 358 | def _get_extensions(folder: str): |
| 359 | folder = files.get_abs_path(folder) |
| 360 | cached = cache.get(_EXTENSIONS_CACHE_AREA, folder) |
| 361 | if cached is not None: |
| 362 | return cached |
| 363 | |
| 364 | if not files.exists(folder): |
| 365 | return [] |
| 366 | |
| 367 | classes = modules.load_classes_from_folder(folder, "*", Extension) |
| 368 | cache.add(_EXTENSIONS_CACHE_AREA, folder, classes) |
| 369 | return classes |
| 370 | |
| 371 | |
| 372 | def register_extensions_watchdogs(): |
| 373 | from helpers import watchdog, projects |
| 374 | |
| 375 | def extensions_changed(items: list[watchdog.WatchItem]): |
| 376 | cache.clear(_EXTENSIONS_CACHE_AREA) |
| 377 | cache.clear(_CLASSES_CACHE_AREA) |
| 378 | PrintStyle.debug("Extensions watchdog triggered:", items) |
| 379 | |
| 380 | # extensions and usr/extensions |
| 381 | watchdog.add_watchdog( |
| 382 | id="extensions_base", |
| 383 | roots=[ |
| 384 | files.get_abs_path(files.EXTENSIONS_DIR), |
| 385 | files.get_abs_path(files.USER_DIR, files.EXTENSIONS_DIR), |
| 386 | ], |
| 387 | handler=extensions_changed, |
| 388 | ) |
| 389 | |
| 390 | # usr/projects/**/extensions |
| 391 | watchdog.add_watchdog( |
| 392 | id="extensions_projects", |
| 393 | roots=[projects.PROJECTS_PARENT_DIR], |
| 394 | patterns=[f"*/{projects.PROJECT_META_DIR}/**/{files.EXTENSIONS_DIR}/**/*"], |
| 395 | handler=extensions_changed, |
| 396 | ) |
| 397 | |
| 398 | # agents and usr/agents |
| 399 | watchdog.add_watchdog( |
| 400 | id="extensions_agents", |
| 401 | roots=[ |
| 402 | files.get_abs_path(files.AGENTS_DIR), |
| 403 | files.get_abs_path(files.USER_DIR, files.AGENTS_DIR), |
| 404 | ], |
| 405 | patterns=[f"*/{files.EXTENSIONS_DIR}/**/*"], |
| 406 | handler=extensions_changed, |
| 407 | ) |