@samitouri / QOSami_Agents_Zero / commits / f7b76839

Add extensible hooks and improve agent detection

Annotate many Agent and AgentContext methods with @extensible to expose extension points and enable runtime hooks. Refactor python/helpers/extension.py to derive extension points from __qualname__, add a _get_agent helper that more robustly finds an Agent instance in args/kwargs, and skip extension handling when module/qualname cannot be determined. Apply extensible to relevant Flask handlers in run_ui and import the extension helper. Also bump the default sleep_time in the code execution tool from 0.1 to 0.5s to reduce polling frequency.

frdel committed Feb 25, 2026 at 20:32 UTC f7b7683918b7148e113d7ffd3e6372b28de96fd6

4 files changed +62 -16

agent.py

+34 -2

@@ -15,7 +15,7 @@ from python.helpers import (

15 tokens,

16 context as context_helper,

17 dirty_json,

18 - subagents

18 + subagents,

19 )

20 from python.helpers.print_style import PrintStyle

@@ -29,7 +29,7 @@ from python.helpers.dirty_json import DirtyJson

29 from python.helpers.defer import DeferredTask

30 from typing import Callable

31 from python.helpers.localization import Localization

32 -from python.helpers.extension import call_extensions

32 +from python.helpers.extension import call_extensions, extensible

33 from python.helpers.errors import RepairableException

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

46 _counter: int = 0

47 _notification_manager = None

49 + @extensible

50 def __init__(

51 self,

52 config: "AgentConfig",

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

152 return cls._notification_manager

153

154 @staticmethod

155 + @extensible

156 def remove(id: str):

157 with AgentContext._contexts_lock:

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

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

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

177 self.output_data[key] = value

178

179 + @extensible

180 def output(self):

181 return {

182 "id": self.id,

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

220 )

221 return items

222

223 + @extensible

224 def kill_process(self):

225 if self.task:

226 self.task.kill()

227

228 + @extensible

229 def reset(self):

230 self.kill_process()

231 self.log.reset()

@@ -228,6 +233,7 @@ class AgentContext:

233 self.streaming_agent = None

234 self.paused = False

235

236 + @extensible

237 def nudge(self):

238 self.kill_process()

239 self.paused = False

@@ -240,6 +246,7 @@ class AgentContext:

246 def is_running(self) -> bool:

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

248

249 + @extensible

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

251 self.paused = False # unpause if paused

252

@@ -259,6 +266,7 @@ class AgentContext:

266

267 return self.task

268

269 + @extensible

270 def run_task(

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

272 ):

@@ -270,6 +278,7 @@ class AgentContext:

278 return self.task

279

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

281 + @extensible

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

283 try:

284 msg_template = (

@@ -355,6 +364,7 @@ class Agent:

364 DATA_NAME_SUBORDINATE = "_subordinate"

365 DATA_NAME_CTX_WINDOW = "ctx_window"

366

367 + @extensible

368 def __init__(

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

370 ):

@@ -376,6 +386,7 @@ class Agent:

386

387 asyncio.run(self.call_extensions("agent_init"))

388

389 + @extensible

390 async def monologue(self):

391 error_retries = 0 # counter for critical error retries

392 while True:

@@ -530,6 +541,7 @@ class Agent:

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

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

543

544 + @extensible

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

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

547

@@ -581,6 +593,7 @@ class Agent:

593

594 return full_prompt

595

596 + @extensible

597 async def retry_critical_exception(

598 self, e: Exception, error_retries: int, delay: int = 3, max_retries: int = 1

599 ) -> int:

@@ -606,6 +619,7 @@ class Agent:

619 )

620 return error_retries + 1

621

622 + @extensible

623 def handle_critical_exception(self, exception: Exception):

624 if isinstance(exception, HandledException):

625 raise exception # Re-raise the exception to kill the loop

@@ -634,6 +648,7 @@ class Agent:

648

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

650

651 + @extensible

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

653 system_prompt: list[str] = []

654 await self.call_extensions(

@@ -641,6 +656,7 @@ class Agent:

656 )

657 return system_prompt

658

659 + @extensible

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

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

662

@@ -649,6 +665,7 @@ class Agent:

665 )

666 return prompt

667

668 + @extensible

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

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

671

@@ -663,6 +680,7 @@ class Agent:

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

681 self.data[field] = value

682

683 + @extensible

684 def hist_add_message(

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

686 ):

@@ -676,6 +694,7 @@ class Agent:

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

695 )

696

697 + @extensible

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

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

700

@@ -704,15 +723,18 @@ class Agent:

723 self.last_user_message = msg

724 return msg

725

726 + @extensible

727 def hist_add_ai_response(self, message: str):

728 self.loop_data.last_response = message

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

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

731

732 + @extensible

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

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

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

736

737 + @extensible

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

739 data = {

740 "tool_name": tool_name,

@@ -727,6 +749,7 @@ class Agent:

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

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

751

752 + @extensible

753 def get_chat_model(self):

754 return models.get_chat_model(

755 self.config.chat_model.provider,

@@ -735,6 +758,7 @@ class Agent:

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

759 )

760

761 + @extensible

762 def get_utility_model(self):

763 return models.get_chat_model(

764 self.config.utility_model.provider,

@@ -743,6 +767,7 @@ class Agent:

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

768 )

769

770 + @extensible

771 def get_browser_model(self):

772 return models.get_browser_model(

773 self.config.browser_model.provider,

@@ -751,6 +776,7 @@ class Agent:

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

777 )

778

779 + @extensible

780 def get_embedding_model(self):

781 return models.get_embedding_model(

782 self.config.embeddings_model.provider,

@@ -759,6 +785,7 @@ class Agent:

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

786 )

787

788 + @extensible

789 async def call_utility_model(

790 self,

791 system: str,

@@ -794,6 +821,7 @@ class Agent:

821

822 return response

823

824 + @extensible

825 async def call_chat_model(

826 self,

827 messages: list[BaseMessage],

@@ -820,6 +848,7 @@ class Agent:

848

849 return response, reasoning

850

851 + @extensible

852 async def rate_limiter_callback(

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

854 ):

@@ -827,6 +856,7 @@ class Agent:

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

857 return False

858

859 + @extensible

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

861 while self.context.paused:

862 await asyncio.sleep(0.1) # wait if paused

@@ -852,6 +882,7 @@ class Agent:

882 while self.context.paused:

883 await asyncio.sleep(0.1)

884

885 + @extensible

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

887 # search for tool usage requests in agent message

888 tool_request = extract_tools.json_parse_dirty(msg)

@@ -971,6 +1002,7 @@ class Agent:

1002 except Exception as e:

1003 pass

1004

1005 + @extensible

1006 def get_tool(

1007 self,

1008 name: str,

python/helpers/extension.py

+21 -12

@@ -64,20 +64,28 @@ def extensible(func):

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"

67 + qual_name = getattr(func, "__qualname__", "")

68 +

69 + # skip if extension point cannot be determined

70 + if not module_name or not qual_name:

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

72 +

73 + start_point = f"{module_name}.{qual_name}-start"

74 + end_point = f"{module_name}.{qual_name}-end"

75 +

76 + def _get_agent() -> "Agent|None":

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

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

79 + return candidate

80 +

81 + for a in args:

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

83 + return a

84 +

85 + return None

87 # 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

88 + agent = _get_agent()

90 # build extension data object - func input/output

91 data = {

@@ -106,6 +114,7 @@ def extensible(func):

114 data["exception"] = e

115

116 # call end extensions, these can modify outputs or exception

117 + agent = _get_agent()

118 await call_extensions(end_point, agent=agent, data=data)

119

120 # if there's an exception, raise it

python/tools/code_execution_tool.py

+1 -1

@@ -241,7 +241,7 @@ class CodeExecution(Tool):

241 between_output_timeout=15, # Wait up to x seconds between outputs

242 dialog_timeout=5, # potential dialog detection timeout

243 max_exec_timeout=180, # hard cap on total runtime

244 - sleep_time=0.1,

244 + sleep_time=0.5,

245 prefix="",

246 timeouts: dict | None = None,

247 ):

run_ui.py

+6 -1

@@ -11,7 +11,7 @@ from flask import Flask, request, Response, session, redirect, url_for, render_t

11 from werkzeug.wrappers.request import Request as WerkzeugRequest

13 import initialize

14 -from python.helpers import files, git, mcp_server, fasta2a_server, settings as settings_helper

14 +from python.helpers import files, git, mcp_server, fasta2a_server, settings as settings_helper, extension

15 from python.helpers.files import get_abs_path

16 from python.helpers import runtime, dotenv, process

17 from python.helpers.websocket import WebSocketHandler, validate_ws_origin

@@ -83,6 +83,7 @@ websocket_manager.set_server_restart_broadcast(

85 @webapp.route("/login", methods=["GET", "POST"])

86 +@extension.extensible

87 async def login_handler():

88 error = None

89 if request.method == 'POST':

@@ -101,6 +102,7 @@ async def login_handler():

102

103

104 @webapp.route("/logout")

105 +@extension.extensible

106 async def logout_handler():

107 session.pop('authentication', None)

108 return redirect(url_for('login_handler'))

@@ -109,6 +111,7 @@ async def logout_handler():

111 # handle default address, load index

112 @webapp.route("/", methods=["GET"])

113 @requires_auth

114 +@extension.extensible

115 async def serve_index():

116 gitinfo = None

117 try:

@@ -142,6 +145,7 @@ async def serve_plugin_asset(plugin_name, asset_path):

145 return await _serve_plugin_asset(plugin_name, asset_path)

146

147

148 +@extension.extensible

149 async def _serve_plugin_asset(plugin_name, asset_path):

150 """

151 Serve static assets from plugin directories.

@@ -452,6 +456,7 @@ def wait_for_health(host: str, port: int):

456 time.sleep(1)

457

458

459 +@extension.extensible

460 def init_a0():

461 # initialize contexts and MCP

462 init_chats = initialize.initialize_chats()