@samitouri / QOSami_Agents_Zero-2 / commits / ab9fc4ee

Refactor extensions to async/sync API

Redesign extension handling to support explicit async/sync execution. helpers/extension.py rewrites the extensible decorator, adds call_extensions_async / call_extensions_sync, and a helper to gather extension classes; caching flag adjusted. Updated call sites across the codebase (agent, APIs, plugins, tools, settings, extensions) to use extension.extensible and the new call_extensions_async/sync API, and converted several extension handlers from async to sync. Also small frontend tweaks (use globalThis.runtimeInfo) and minor import updates (csrf_protect in run_ui). This centralizes extension discovery/execution and avoids previously scattered asyncio.run usage.

frdel committed Mar 6, 2026 at 11:32 UTC ab9fc4ee7f3d011069a6027b717c6d7db61a897c

14 files changed +237 -182

agent.py

+82 -69

@@ -30,7 +30,7 @@ from helpers.dirty_json import DirtyJson

30 from helpers.defer import DeferredTask

31 from typing import Callable

32 from helpers.localization import Localization

33 -from helpers.extension import call_extensions, extensible

33 +from helpers import extension

34 from helpers.errors import RepairableException, InterventionException, HandledException

36 class AgentContextType(Enum):

@@ -46,7 +46,7 @@ class AgentContext:

46 _counter: int = 0

47 _notification_manager = None

49 - @extensible

49 + @extension.extensible

50 def __init__(

51 self,

52 config: "AgentConfig",

@@ -152,7 +152,7 @@ class AgentContext:

152 return cls._notification_manager

153

154 @staticmethod

155 - @extensible

155 + @extension.extensible

156 def remove(id: str):

157 with AgentContext._contexts_lock:

158 context = AgentContext._contexts.pop(id, None)

@@ -176,7 +176,7 @@ class AgentContext:

176 # recursive is not used now, prepared for context hierarchy

177 self.output_data[key] = value

178

179 - # @extensible

179 + @extension.extensible

180 def output(self):

181 return {

182 "id": self.id,

@@ -220,12 +220,12 @@ class AgentContext:

220 )

221 return items

222

223 - @extensible

223 + @extension.extensible

224 def kill_process(self):

225 if self.task:

226 self.task.kill()

227

228 - @extensible

228 + @extension.extensible

229 def reset(self):

230 self.kill_process()

231 self.log.reset()

@@ -233,21 +233,21 @@ class AgentContext:

233 self.streaming_agent = None

234 self.paused = False

235

236 - @extensible

236 + @extension.extensible

237 def nudge(self):

238 self.kill_process()

239 self.paused = False

240 self.task = self.communicate(UserMessage(self.agent0.read_prompt("fw.msg_nudge.md")))

241 return self.task

242

243 - @extensible

243 + @extension.extensible

244 def get_agent(self):

245 return self.streaming_agent or self.agent0

246

247 def is_running(self) -> bool:

248 return (self.task and self.task.is_alive()) or False

249

250 - @extensible

250 + @extension.extensible

251 def communicate(self, msg: "UserMessage", broadcast_level: int = 1):

252 self.paused = False # unpause if paused

253

@@ -267,7 +267,7 @@ class AgentContext:

267

268 return self.task

269

270 - @extensible

270 + @extension.extensible

271 def run_task(

272 self, func: Callable[..., Coroutine[Any, Any, Any]], *args: Any, **kwargs: Any

273 ):

@@ -279,7 +279,7 @@ class AgentContext:

279 return self.task

280

281 # this wrapper ensures that superior agents are called back if the chat was loaded from file and original callstack is gone

282 - @extensible

282 + @extension.extensible

283 async def _process_chain(self, agent: "Agent", msg: "UserMessage|str", user=True):

284 try:

285 msg_template = (

@@ -295,13 +295,13 @@ class AgentContext:

295 response = await self._process_chain(superior, response, False) # type: ignore

296

297 # call end of process extensions

298 - await self.get_agent().call_extensions("process_chain_end", data={})

298 + await extension.call_extensions_async("process_chain_end", agent=self.get_agent(), data={})

299

300 return response

301 except Exception as e:

302 await self.handle_exception("process_chain", e)

303

304 - @extensible

304 + @extension.extensible

305 async def handle_exception(self, location: str, exception: Exception):

306 if exception:

307 raise exception # exception handling is done by extensions

@@ -358,7 +358,7 @@ class Agent:

358 DATA_NAME_SUBORDINATE = "_subordinate"

359 DATA_NAME_CTX_WINDOW = "ctx_window"

360

361 - @extensible

361 + @extension.extensible

362 def __init__(

363 self, number: int, config: AgentConfig, context: AgentContext | None = None

364 ):

@@ -378,16 +378,18 @@ class Agent:

378 self.intervention: UserMessage | None = None

379 self.data: dict[str, Any] = {} # free data object all the tools can use

380

381 - asyncio.run(self.call_extensions("agent_init"))

381 + extension.call_extensions_sync("agent_init", self)

382

383 - @extensible

383 + @extension.extensible

384 async def monologue(self):

385 while True:

386 try:

387 # loop data dictionary to pass to extensions

388 self.loop_data = LoopData(user_message=self.last_user_message)

389 # call monologue_start extensions

390 - await self.call_extensions("monologue_start", loop_data=self.loop_data)

390 + await extension.call_extensions_async(

391 + "monologue_start", self, loop_data=self.loop_data

392 + )

393

394 printer = PrintStyle(italic=True, font_color="#b3ffd9", padding=False)

395

@@ -399,8 +401,8 @@ class Agent:

401 self.loop_data.params_temporary = {} # clear temporary params

402

403 # call message_loop_start extensions

402 - await self.call_extensions(

403 - "message_loop_start", loop_data=self.loop_data

404 + await extension.call_extensions_async(

405 + "message_loop_start", self, loop_data=self.loop_data

406 )

407 await self.handle_intervention()

408

@@ -409,8 +411,8 @@ class Agent:

411 prompt = await self.prepare_prompt(loop_data=self.loop_data)

412

413 # call before_main_llm_call extensions

412 - await self.call_extensions(

413 - "before_main_llm_call", loop_data=self.loop_data

414 + await extension.call_extensions_async(

415 + "before_main_llm_call", self, loop_data=self.loop_data

416 )

417 await self.handle_intervention()

418

@@ -421,8 +423,9 @@ class Agent:

423 printer.print("Reasoning: ") # start of reasoning

424 # Pass chunk and full data to extensions for processing

425 stream_data = {"chunk": chunk, "full": full}

424 - await self.call_extensions(

426 + await extension.call_extensions_async(

427 "reasoning_stream_chunk",

428 + self,

429 loop_data=self.loop_data,

430 stream_data=stream_data,

431 )

@@ -439,8 +442,9 @@ class Agent:

442 printer.print("Response: ") # start of response

443 # Pass chunk and full data to extensions for processing

444 stream_data = {"chunk": chunk, "full": full}

442 - await self.call_extensions(

445 + await extension.call_extensions_async(

446 "response_stream_chunk",

447 + self,

448 loop_data=self.loop_data,

449 stream_data=stream_data,

450 )

@@ -459,13 +463,13 @@ class Agent:

463 await self.handle_intervention(agent_response)

464

465 # Notify extensions to finalize their stream filters

462 - await self.call_extensions(

463 - "reasoning_stream_end", loop_data=self.loop_data

466 + await extension.call_extensions_async(

467 + "reasoning_stream_end", self, loop_data=self.loop_data

468 )

469 await self.handle_intervention(agent_response)

470

467 - await self.call_extensions(

468 - "response_stream_end", loop_data=self.loop_data

471 + await extension.call_extensions_async(

472 + "response_stream_end", self, loop_data=self.loop_data

473 )

474

475 await self.handle_intervention(agent_response)

@@ -498,8 +502,8 @@ class Agent:

502 finally:

503 # call message_loop_end extensions

504 if self.context.task and self.context.task.is_alive(): # don't call extensions post mortem

501 - await self.call_extensions(

502 - "message_loop_end", loop_data=self.loop_data

505 + await extension.call_extensions_async(

506 + "message_loop_end", self, loop_data=self.loop_data

507 )

508

509

@@ -511,21 +515,27 @@ class Agent:

515 self.context.streaming_agent = None # unset current streamer

516 # call monologue_end extensions

517 if self.context.task and self.context.task.is_alive(): # don't call extensions post mortem

514 - await self.call_extensions("monologue_end", loop_data=self.loop_data) # type: ignore

518 + await extension.call_extensions_async(

519 + "monologue_end", self, loop_data=self.loop_data

520 + ) # type: ignore

521

516 - @extensible

522 + @extension.extensible

523 async def prepare_prompt(self, loop_data: LoopData) -> list[BaseMessage]:

524 self.context.log.set_progress("Building prompt")

525

526 # call extensions before setting prompts

521 - await self.call_extensions("message_loop_prompts_before", loop_data=loop_data)

527 + await extension.call_extensions_async(

528 + "message_loop_prompts_before", self, loop_data=loop_data

529 + )

530

531 # set system prompt and message history

532 loop_data.system = await self.get_system_prompt(self.loop_data)

533 loop_data.history_output = self.history.output()

534

535 # and allow extensions to edit them

528 - await self.call_extensions("message_loop_prompts_after", loop_data=loop_data)

536 + await extension.call_extensions_async(

537 + "message_loop_prompts_after", self, loop_data=loop_data

538 + )

539

540 # concatenate system prompt

541 system_text = "\n\n".join(loop_data.system)

@@ -565,7 +575,7 @@ class Agent:

575

576 return full_prompt

577

568 - @extensible

578 + @extension.extensible

579 async def handle_exception(self, location: str, exception: Exception):

580 if exception:

581 raise exception # exception handling is done by extensions

@@ -605,15 +615,15 @@ class Agent:

615

616 # raise HandledException(exception) # Re-raise the exception to kill the loop

617

608 - @extensible

618 + @extension.extensible

619 async def get_system_prompt(self, loop_data: LoopData) -> list[str]:

620 system_prompt: list[str] = []

611 - await self.call_extensions(

612 - "system_prompt", system_prompt=system_prompt, loop_data=loop_data

621 + await extension.call_extensions_async(

622 + "system_prompt", self, system_prompt=system_prompt, loop_data=loop_data

623 )

624 return system_prompt

625

616 - @extensible

626 + @extension.extensible

627 def parse_prompt(self, _prompt_file: str, **kwargs):

628 dirs = subagents.get_paths(self, "prompts")

629

@@ -622,7 +632,7 @@ class Agent:

632 )

633 return prompt

634

625 - @extensible

635 + @extension.extensible

636 def read_prompt(self, file: str, **kwargs) -> str:

637 dirs = subagents.get_paths(self, "prompts")

638

@@ -637,21 +647,21 @@ class Agent:

647 def set_data(self, field: str, value):

648 self.data[field] = value

649

640 - @extensible

650 + @extension.extensible

651 def hist_add_message(

652 self, ai: bool, content: history.MessageContent, tokens: int = 0

653 ):

654 self.last_message = datetime.now(timezone.utc)

655 # Allow extensions to process content before adding to history

656 content_data = {"content": content}

647 - asyncio.run(

648 - self.call_extensions("hist_add_before", content_data=content_data, ai=ai)

657 + extension.call_extensions_sync(

658 + "hist_add_before", self, content_data=content_data, ai=ai

659 )

660 return self.history.add_message(

661 ai=ai, content=content_data["content"], tokens=tokens

662 )

663

654 - @extensible

664 + @extension.extensible

665 def hist_add_user_message(self, message: UserMessage, intervention: bool = False):

666 self.history.new_topic() # user message starts a new topic in history

667

@@ -680,25 +690,25 @@ class Agent:

690 self.last_user_message = msg

691 return msg

692

683 - @extensible

693 + @extension.extensible

694 def hist_add_ai_response(self, message: str):

695 self.loop_data.last_response = message

696 content = self.parse_prompt("fw.ai_response.md", message=message)

697 return self.hist_add_message(True, content=content)

698

689 - @extensible

699 + @extension.extensible

700 def hist_add_warning(self, message: history.MessageContent):

701 content = self.parse_prompt("fw.warning.md", message=message)

702 return self.hist_add_message(False, content=content)

703

694 - @extensible

704 + @extension.extensible

705 def hist_add_tool_result(self, tool_name: str, tool_result: str, **kwargs):

706 data = {

707 "tool_name": tool_name,

708 "tool_result": tool_result,

709 **kwargs,

710 }

701 - asyncio.run(self.call_extensions("hist_add_tool_result", data=data))

711 + extension.call_extensions_sync("hist_add_tool_result", self, data=data)

712 return self.hist_add_message(False, content=data)

713

714 def concat_messages(

@@ -706,7 +716,7 @@ class Agent:

716 ): # TODO add param for message range, topic, history

717 return self.history.output_text(human_label="user", ai_label="assistant")

718

709 - @extensible

719 + @extension.extensible

720 def get_chat_model(self):

721 return models.get_chat_model(

722 self.config.chat_model.provider,

@@ -715,7 +725,7 @@ class Agent:

725 **self.config.chat_model.build_kwargs(),

726 )

727

718 - @extensible

728 + @extension.extensible

729 def get_utility_model(self):

730 return models.get_chat_model(

731 self.config.utility_model.provider,

@@ -724,7 +734,7 @@ class Agent:

734 **self.config.utility_model.build_kwargs(),

735 )

736

727 - @extensible

737 + @extension.extensible

738 def get_browser_model(self):

739 return models.get_browser_model(

740 self.config.browser_model.provider,

@@ -733,7 +743,7 @@ class Agent:

743 **self.config.browser_model.build_kwargs(),

744 )

745

736 - @extensible

746 + @extension.extensible

747 def get_embedding_model(self):

748 return models.get_embedding_model(

749 self.config.embeddings_model.provider,

@@ -742,7 +752,7 @@ class Agent:

752 **self.config.embeddings_model.build_kwargs(),

753 )

754

745 - @extensible

755 + @extension.extensible

756 async def call_utility_model(

757 self,

758 system: str,

@@ -760,7 +770,9 @@ class Agent:

770 "callback": callback,

771 "background": background,

772 }

763 - await self.call_extensions("util_model_call_before", call_data=call_data)

773 + await extension.call_extensions_async(

774 + "util_model_call_before", self, call_data=call_data

775 + )

776

777 # propagate stream to callback if set

778 async def stream_callback(chunk: str, total: str):

@@ -778,7 +790,7 @@ class Agent:

790

791 return response

792

781 - @extensible

793 + @extension.extensible

794 async def call_chat_model(

795 self,

796 messages: list[BaseMessage],

@@ -805,7 +817,7 @@ class Agent:

817

818 return response, reasoning

819

808 - @extensible

820 + @extension.extensible

821 async def rate_limiter_callback(

822 self, message: str, key: str, total: int, limit: int

823 ):

@@ -813,7 +825,7 @@ class Agent:

825 self.context.log.set_progress(message, True)

826 return False

827

816 - @extensible

828 + @extension.extensible

829 async def handle_intervention(self, progress: str = ""):

830 await self.wait_if_paused()

831 if (

@@ -838,7 +850,7 @@ class Agent:

850 while self.context.paused:

851 await asyncio.sleep(0.1)

852

841 - @extensible

853 + @extension.extensible

854 async def process_tools(self, msg: str):

855 # search for tool usage requests in agent message

856 tool_request = extract_tools.json_parse_dirty(msg)

@@ -894,8 +906,9 @@ class Agent:

906 await self.handle_intervention()

907

908 # Allow extensions to preprocess tool arguments

897 - await self.call_extensions(

909 + await extension.call_extensions_async(

910 "tool_execute_before",

911 + self,

912 tool_args=tool_args or {},

913 tool_name=tool_name,

914 )

@@ -904,8 +917,11 @@ class Agent:

917 await self.handle_intervention()

918

919 # Allow extensions to postprocess tool response

907 - await self.call_extensions(

908 - "tool_execute_after", response=response, tool_name=tool_name

920 + await extension.call_extensions_async(

921 + "tool_execute_after",

922 + self,

923 + response=response,

924 + tool_name=tool_name,

925 )

926

927 await tool.after_execution(response)

@@ -935,8 +951,9 @@ class Agent:

951

952 async def handle_reasoning_stream(self, stream: str):

953 await self.handle_intervention()

938 - await self.call_extensions(

954 + await extension.call_extensions_async(

955 "reasoning_stream",

956 + self,

957 loop_data=self.loop_data,

958 text=stream,

959 )

@@ -948,8 +965,9 @@ class Agent:

965 return # no reason to try

966 response = DirtyJson.parse_string(stream)

967 if isinstance(response, dict):

951 - await self.call_extensions(

968 + await extension.call_extensions_async(

969 "response_stream",

970 + self,

971 loop_data=self.loop_data,

972 text=stream,

973 parsed=response,

@@ -958,7 +976,7 @@ class Agent:

976 except Exception as e:

977 pass

978

961 - @extensible

979 + @extension.extensible

980 def get_tool(

981 self,

982 name: str,

@@ -992,9 +1010,4 @@ class Agent:

1010 message=message,

1011 loop_data=loop_data,

1012 **kwargs,

995 - )

996 -

997 - async def call_extensions(self, extension_point: str, **kwargs) -> Any:

998 - return await call_extensions(

999 - extension_point=extension_point, agent=self, **kwargs

1000 - )

1013 + )

\ No newline at end of file

api/banners.py

+2 -2

@@ -1,5 +1,5 @@

1 from helpers.api import ApiHandler, Request, Response

2 -from helpers.extension import call_extensions

2 +from helpers.extension import call_extensions_async

5 class GetBanners(ApiHandler):

@@ -13,7 +13,7 @@ class GetBanners(ApiHandler):

13 frontend_context = input.get("context", {})

15 # Banners array passed by reference - extensions append directly to it

16 - await call_extensions("banners", agent=None, banners=banners, frontend_context=frontend_context)

16 + await call_extensions_async("banners", agent=None, banners=banners, frontend_context=frontend_context)

18 return {"banners": banners}

api/message.py

+1 -1

@@ -58,7 +58,7 @@ class Message(ApiHandler):

59 # call extension point, alow it to modify data

60 data = { "message": message, "attachment_paths": attachment_paths }

61 - await extension.call_extensions("user_message_ui", agent=context.get_agent(), data=data)

61 + await extension.call_extensions_async("user_message_ui", agent=context.get_agent(), data=data)

62 message = data.get("message", "")

63 attachment_paths = data.get("attachment_paths", [])

extensions/python/agent_Agent_handle_exception_end/_50_handle_repairable_exception.py

+2 -3

@@ -3,10 +3,9 @@ from helpers.extension import Extension

3 from agent import LoopData

4 from helpers.localization import Localization

5 from helpers.errors import RepairableException

6 -from helpers import errors

6 +from helpers import errors, extension

7 from helpers.print_style import PrintStyle

9 -

9 class HandleRepairableException(Extension):

10 async def execute(self, data: dict = {}, **kwargs):

11 if not self.agent:

@@ -17,7 +16,7 @@ class HandleRepairableException(Extension):

17 if isinstance(data["exception"], RepairableException):

18 msg = {"message": errors.format_error(data["exception"])}

20 - await self.agent.call_extensions("error_format", msg=msg)

19 + await extension.call_extensions_async("error_format", agent=self.agent, msg=msg)

20 self.agent.hist_add_warning(msg["message"])

21 PrintStyle(font_color="red", padding=True).print(msg["message"])

22 self.agent.context.log.log(type="warning", content=msg["message"])

extensions/python/agent_init/_10_initial_message.py

+1 -1

@@ -5,7 +5,7 @@ from helpers.extension import Extension

6 class InitialMessage(Extension):

8 - async def execute(self, **kwargs):

8 + def execute(self, **kwargs):

9 """

10 Add an initial greeting message when first user message is processed.

11 Called only once per session via _process_chain method.

extensions/python/agent_init/_15_load_profile_settings.py

+1 -1

@@ -5,7 +5,7 @@ from helpers.extension import Extension

6 class LoadProfileSettings(Extension):

8 - async def execute(self, **kwargs) -> None:

8 + def execute(self, **kwargs) -> None:

10 if not self.agent or not self.agent.config.profile:

11 return

extensions/python/hist_add_before/_10_mask_content.py

+1 -1

@@ -4,7 +4,7 @@ from helpers.secrets import get_secrets_manager

5 class MaskHistoryContent(Extension):

7 - async def execute(self, **kwargs):

7 + def execute(self, **kwargs):

8 if not self.agent:

9 return

extensions/python/hist_add_tool_result/_90_save_tool_call_file.py

+1 -1

@@ -6,7 +6,7 @@ import os, re

6 LEN_MIN = 500

8 class SaveToolCallFile(Extension):

9 - async def execute(self, data: dict[str, Any] | None = None, **kwargs):

9 + def execute(self, data: dict[str, Any] | None = None, **kwargs):

10 if not self.agent:

11 return

helpers/extension.py

+125 -83

@@ -1,10 +1,9 @@

1 from abc import abstractmethod

2 -from typing import Any

2 +from typing import Any, Awaitable, Type, cast

3 from helpers import extract_tools, files

4 from helpers import cache, plugins, subagents

5 from typing import TYPE_CHECKING

6 from functools import wraps

7 -import asyncio

7 import inspect

9 if TYPE_CHECKING:

@@ -15,7 +14,7 @@ DEFAULT_EXTENSIONS_FOLDER = "python/extensions"

14 USER_EXTENSIONS_FOLDER = "usr/extensions"

16 _CACHE_AREA = "extension_folder_classes(extensions)(plugins)"

18 -cache.toggle_area(_CACHE_AREA, True) # cache off for now

17 +cache.toggle_area(_CACHE_AREA, False) # cache off for now

20 class _Unset:

@@ -35,62 +34,54 @@ def extensible(func):

34 - ``{func.__module__}_{func.__qualname__}_end`` with `.` replaced by `_`

36 When the wrapped function is called, the decorator builds a mutable ``data``

38 - payload and passes it to both extension points via ``call_extensions``:

37 + payload and passes it to both extension points:

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

39 + - ``data["args"]``: positional args (extensions may replace/mutate)

40 + - ``data["kwargs"]``: keyword args (extensions may replace/mutate)

41 + - ``data["result"]``: initialized to an internal sentinel; extensions may set

42 + this to short-circuit the wrapped function

43 - ``data["exception"]``: initialized to an internal sentinel; extensions may

44 set this to a ``BaseException`` instance to force-raise

46 + Sync functions call ``call_extensions_sync``. Async functions call

47 + ``call_extensions_async``.

48 +

49 Behavior:

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"]``.

51 + - ``-start`` extensions run first and may mutate inputs or set

52 + ``data["result"]`` / ``data["exception"]``.

53 + - If ``data["result"]`` is still unset, the decorator calls the wrapped

54 + function using the possibly modified ``data["args"]`` / ``data["kwargs"]``.

55 + - ``-end`` extensions run last and may rewrite ``data["result"]`` or replace /

56 + clear ``data["exception"]``.

58 Finally, if ``data["exception"]`` contains an exception it is raised;

59 otherwise ``data["result"]`` is returned.

60 """

62 - @wraps(func)

63 - async def _inner_async(*args, **kwargs):

62 + def _get_agent(args, kwargs):

63 from agent import Agent

66 - # prepare extension points data

65 + candidate = kwargs.get("agent")

66 + if isinstance(candidate, Agent) and bool(getattr(candidate, "__dict__", None)):

67 + return candidate

68 +

69 + for a in args:

70 + if isinstance(a, Agent) and bool(getattr(a, "__dict__", None)):

71 + return a

72 +

73 + return None

74 +

75 + def _prepare_inputs(args, kwargs):

76 module_name = getattr(func, "__module__", "").replace(".", "_")

77 qual_name = getattr(func, "__qualname__", "").replace(".", "_")

69 -

70 - # skip if extension point cannot be determined

78 if not module_name or not qual_name:

72 - return await func(*args, **kwargs)

79 + return None

81 start_point = f"{module_name}_{qual_name}_start"

82 end_point = f"{module_name}_{qual_name}_end"

83 + agent = _get_agent(args, kwargs)

77 - def _get_agent() -> "Agent|None":

78 - candidate = kwargs.get("agent")

79 - if isinstance(candidate, Agent) and bool(

80 - getattr(candidate, "__dict__", None)

81 - ):

82 - return candidate

83 -

84 - for a in args:

85 - if isinstance(a, Agent) and bool(getattr(a, "__dict__", None)):

86 - return a

87 -

88 - return None

89 -

90 - # try to find agent instance for better extension determination

91 - agent = _get_agent()

92 -

93 - # build extension data object - func input/output

85 data = {

86 "args": args,

87 "kwargs": kwargs,

@@ -98,44 +89,78 @@ def extensible(func):

89 "exception": None,

90 }

101 - # call start extensions, these can modify inputs, produce output or exception

102 - await call_extensions(start_point, agent=agent, data=data)

92 + return start_point, end_point, agent, data

104 - # if there is an explicit exception set, raise it

94 + def _process_result(data):

95 exc = data.get("exception")

96 if isinstance(exc, BaseException):

97 raise exc

109 - # if there is no result set, call the original function

110 - if data.get("result") is _UNSET:

99 + return data.get("result")

100 +

101 + def _call_original(data):

102 + call_args = data.get("args")

103 + call_kwargs = data.get("kwargs")

104 +

105 + if not isinstance(call_args, tuple):

106 + call_args = (call_args,)

107 + if not isinstance(call_kwargs, dict):

108 + call_kwargs = {}

109 +

110 + try:

111 + data["result"] = func(*call_args, **call_kwargs)

112 + except Exception as e:

113 + data["exception"] = e

114 + return _UNSET

115 +

116 + async def _run_async(*args, **kwargs):

117 + prepared = _prepare_inputs(args, kwargs)

118 + if prepared is None:

119 + return await func(*args, **kwargs)

120 +

121 + start_point, end_point, agent, data = prepared

122 +

123 + # call pre-extensions

124 + await call_extensions_async(start_point, agent=agent, data=data)

125 +

126 + # call the original if pre-extensions don't return a result

127 + if (result := _process_result(data)) is _UNSET:

128 + _call_original(data)

129 try:

112 - if inspect.iscoroutinefunction(func):

113 - data["result"] = await func(*args, **kwargs)

114 - else:

115 - data["result"] = func(*args, **kwargs)

130 + data["result"] = await data["result"]

131 except Exception as e:

132 data["exception"] = e

133

119 - # call end extensions, these can modify outputs or exception

120 - await call_extensions(end_point, agent=agent, data=data)

134 + # call post-extensions

135 + await call_extensions_async(end_point, agent=agent, data=data)

136

122 - # if there's an exception, raise it

123 - exc = data.get("exception")

124 - if isinstance(exc, BaseException):

125 - raise exc

137 + result = _process_result(data)

138 + return None if result is _UNSET else result

139

127 - # if there's a result, return it

128 - result = data.get("result")

140 + def _run_sync(*args, **kwargs):

141 + prepared = _prepare_inputs(args, kwargs)

142 + if prepared is None:

143 + return func(*args, **kwargs)

144 +

145 + start_point, end_point, agent, data = prepared

146 +

147 + # call pre-extensions

148 + call_extensions_sync(start_point, agent=agent, data=data)

149 +

150 + # call the original if pre-extensions don't return a result

151 + if (result := _process_result(data)) is _UNSET:

152 + _call_original(data)

153 +

154 + # call post-extensions

155 + call_extensions_sync(end_point, agent=agent, data=data)

156 +

157 + result = _process_result(data)

158 return None if result is _UNSET else result

159

160 if inspect.iscoroutinefunction(func):

132 - return _inner_async

133 -

134 - @wraps(func)

135 - def _inner_sync(*args, **kwargs):

136 - return asyncio.run(_inner_async(*args, **kwargs))

161 + return wraps(func)(_run_async)

162

138 - return _inner_sync

163 + return wraps(func)(_run_sync)

164

165

166 class Extension:

@@ -145,37 +170,34 @@ class Extension:

170 self.kwargs = kwargs

171

172 @abstractmethod

148 - async def execute(self, **kwargs) -> Any:

173 + def execute(self, **kwargs) -> None | Awaitable[None]:

174 pass

175

176

152 -async def call_extensions(

177 +async def call_extensions_async(

178 extension_point: str, agent: "Agent|None" = None, **kwargs

154 -) -> Any:

155 - # search for extension folders in all agent's paths

156 - paths = subagents.get_paths(agent, "extensions/python", extension_point)

179 +):

180 + # fetch classes for this extension point and agent

181 + classes = _get_extension_classes(extension_point, agent=agent, **kwargs)

182

158 - # # Add plugin backend extension paths (plugins/*/extensions/python/{extension_point})

159 - # plugin_paths = plugins.get_enabled_plugin_paths(

160 - # agent, "extensions", "python", extension_point

161 - # )

162 - # paths.extend(p for p in plugin_paths if p not in paths)

183 + # execute unique extensions

184 + for cls in classes:

185 + result = cls(agent=agent).execute(**kwargs)

186 + if isinstance(result, Awaitable):

187 + await result

188

164 - all_exts = [cls for path in paths for cls in _get_extensions(path)]

189

166 - # merge: first ocurrence of file name is the override

167 - unique = {}

168 - for cls in all_exts:

169 - file = _get_file_from_module(cls.__module__)

170 - if file not in unique:

171 - unique[file] = cls

172 - classes = sorted(

173 - unique.values(), key=lambda cls: _get_file_from_module(cls.__module__)

174 - )

190 +def call_extensions_sync(extension_point: str, agent: "Agent|None" = None, **kwargs):

191 + # fetch classes for this extension point and agent

192 + classes = _get_extension_classes(extension_point, agent=agent, **kwargs)

193

194 # execute unique extensions

195 for cls in classes:

178 - await cls(agent=agent).execute(**kwargs)

196 + result = cls(agent=agent).execute(**kwargs)

197 + if isinstance(result, Awaitable):

198 + raise ValueError(

199 + f"Extension {cls.__name__} returned awaitable in sync mode"

200 + )

201

202

203 def get_webui_extensions(

@@ -205,6 +227,26 @@ def get_webui_extensions(

227 return entries

228

229

230 +def _get_extension_classes(

231 + extension_point: str, agent: "Agent|None" = None, **kwargs

232 +) -> list[Type[Extension]]:

233 + # search for extension folders in all agent's paths

234 + paths = subagents.get_paths(agent, "extensions/python", extension_point)

235 +

236 + all_exts = [cls for path in paths for cls in _get_extensions(path)]

237 +

238 + # merge: first ocurrence of file name is the override

239 + unique = {}

240 + for cls in all_exts:

241 + file = _get_file_from_module(cls.__module__)

242 + if file not in unique:

243 + unique[file] = cls

244 + classes = sorted(

245 + unique.values(), key=lambda cls: _get_file_from_module(cls.__module__)

246 + )

247 + return classes

248 +

249 +

250 def _get_file_from_module(module_name: str) -> str:

251 return module_name.split(".")[-1]

252

helpers/settings.py

+2 -2

@@ -600,10 +600,10 @@ def _apply_settings(previous: Settings | None):

600 or _settings["embed_model_provider"] != previous["embed_model_provider"]

601 or _settings["embed_model_kwargs"] != previous["embed_model_kwargs"]

602 ):

603 - from helpers.extension import call_extensions

603 + from helpers.extension import call_extensions_async

604

605 defer.DeferredTask().start_task(

606 - call_extensions, "embedding_model_changed"

606 + call_extensions_async, "embedding_model_changed"

607 )

608

609 # update mcp settings if necessary

plugins/text_editor/tools/text_editor.py

+6 -6

@@ -1,5 +1,5 @@

1 from helpers.tool import Tool, Response

2 -from helpers.extension import call_extensions

2 +from helpers.extension import call_extensions_async

3 from helpers import plugins, runtime

4 from plugins.text_editor.helpers.file_ops import (

5 FileInfo,

@@ -63,7 +63,7 @@ class TextEditor(Tool):

63 "content": result["content"],

64 "warnings": result["warnings"],

65 }

66 - await call_extensions(

66 + await call_extensions_async(

67 "text_editor_read_after", agent=self.agent, data=ext_data

68 )

@@ -87,7 +87,7 @@ class TextEditor(Tool):

88 # Extension point

89 ext_data = {"path": path, "content": content}

90 - await call_extensions(

90 + await call_extensions_async(

91 "text_editor_write_before", agent=self.agent, data=ext_data

92 )

@@ -99,7 +99,7 @@ class TextEditor(Tool):

99 return self._error("write", path, result["error"])

100

101 # Extension point

102 - await call_extensions(

102 + await call_extensions_async(

103 "text_editor_write_after", agent=self.agent,

104 data={"path": path, "total_lines": result["total_lines"]},

105 )

@@ -148,7 +148,7 @@ class TextEditor(Tool):

148

149 # Extension point

150 ext_data = {"path": expanded, "edits": parsed}

151 - await call_extensions(

151 + await call_extensions_async(

152 "text_editor_patch_before", agent=self.agent, data=ext_data

153 )

154

@@ -160,7 +160,7 @@ class TextEditor(Tool):

160 return self._error("patch", path, str(exc))

161

162 # Extension point

163 - await call_extensions(

163 + await call_extensions_async(

164 "text_editor_patch_after", agent=self.agent,

165 data={"path": expanded, "total_lines": total_lines},

166 )

run_ui.py

+1 -1

@@ -15,7 +15,7 @@ from helpers import files, git, mcp_server, fasta2a_server, settings as settings

15 from helpers.files import get_abs_path

16 from helpers import runtime, dotenv, process

17 from helpers.websocket import WebSocketHandler, validate_ws_origin

18 -from helpers.api import register_api_route, requires_auth

18 +from helpers.api import register_api_route, requires_auth, csrf_protect

19 from helpers.print_style import PrintStyle

20 from helpers import login

21 import socketio # type: ignore[import-untyped]

webui/js/api.js

+11 -10

@@ -126,10 +126,10 @@ const CSRF_SLOW_WARN_MS = 1500;

126 export function getRuntimeId() {

127 if (runtimeIdCache) return runtimeIdCache;

128 const injected =

129 - window.runtimeInfo &&

130 - typeof window.runtimeInfo.id === "string" &&

131 - window.runtimeInfo.id.length > 0

132 - ? window.runtimeInfo.id

129 + globalThis.runtimeInfo &&

130 + typeof globalThis.runtimeInfo.id === "string" &&

131 + globalThis.runtimeInfo.id.length > 0

132 + ? globalThis.runtimeInfo.id

133 : null;

134 return injected;

135 }

@@ -167,6 +167,7 @@ export async function getCsrfToken() {

167 });

168 }

169

170 + /** @type {RequestInit} */

171 const fetchOptions = { credentials: "same-origin" };

172 if (controller) {

173 fetchOptions.signal = controller.signal;

@@ -177,7 +178,7 @@ export async function getCsrfToken() {

178 ? await Promise.race([fetchPromise, timeoutPromise])

179 : await fetchPromise;

180 } catch (error) {

180 - if (error && error.name === "AbortError") {

181 + if (error && error["name"] === "AbortError") {

182 throw new Error("CSRF token request timed out");

183 }

184 throw error;

@@ -204,10 +205,10 @@ export async function getCsrfToken() {

205 runtimeIdCache = runtimeId;

206 }

207 const injectedRuntimeId =

207 - window.runtimeInfo &&

208 - typeof window.runtimeInfo.id === "string" &&

209 - window.runtimeInfo.id.length > 0

210 - ? window.runtimeInfo.id

208 + globalThis.runtimeInfo &&

209 + typeof globalThis.runtimeInfo.id === "string" &&

210 + globalThis.runtimeInfo.id.length > 0

211 + ? globalThis.runtimeInfo.id

212 : null;

213 const cookieRuntimeId = runtimeId || injectedRuntimeId;

214 if (cookieRuntimeId) {

@@ -216,7 +217,7 @@ export async function getCsrfToken() {

217 console.warn("CSRF runtime id missing; skipping cookie name binding.");

218 }

219 const elapsedMs = Date.now() - startedAt;

219 - if (elapsedMs > CSRF_SLOW_WARN_MS && window.runtimeInfo?.isDevelopment) {

220 + if (elapsedMs > CSRF_SLOW_WARN_MS && globalThis.runtimeInfo?.isDevelopment) {

221 console.warn(`CSRF token request took ${elapsedMs}ms`);

222 }

223 return csrfToken;

webui/js/confirmClick.js

+1 -1

@@ -68,7 +68,7 @@ function resetButton(button, state) {

68 // Register Alpine magic helper

69 export function registerAlpineMagic() {

70 if (globalThis.Alpine) {

71 - Alpine.magic('confirmClick', () => confirmClick);

71 + globalThis.Alpine.magic('confirmClick', () => confirmClick);

72 }

73 }