feat: add tool_output_update extension point for output secrets masking
Adds a new extension hook 'tool_output_update' that fires before any tool output reaches a display or log channel. Changes: - set_progress() promoted to async def — fires hook before updating self.progress - before_execution() arg display loop — fires hook per arg before PrintStyle.stream - Both pass mutable ctx={content:...} so subscribers can rewrite in-place This closes the timing gap where secrets are unmasked between tool_execute_before and tool_execute_after, causing leaks in code_execution output windows and logs. Note: callers of set_progress() must now await it.
Deimos AI committed
Mar 22, 2026 at 11:29 UTC
29197652a27605a148c8891aa626a60762c5bd57
1 file changed
+9
-3
helpers/tool.py
+9
-3
@@ -3,6 +3,7 @@ from dataclasses import dataclass
3
from typing import Any
4
5
from agent import Agent, LoopData
6
+from helpers.extension import call_extensions_async
7
from helpers.print_style import PrintStyle
8
from helpers.strings import sanitize_string
9
@@ -28,8 +29,10 @@ class Tool:
29
async def execute(self,**kwargs) -> Response:
30
pass
31
31
- def set_progress(self, content: str | None):
32
- self.progress = content or ""
32
+ async def set_progress(self, content: str | None):
33
+ ctx = {"content": content or ""}
34
+ await call_extensions_async("tool_output_update", self.agent, ctx=ctx)
35
+ self.progress = ctx["content"]
36
37
def add_progress(self, content: str | None):
38
if not content:
@@ -41,8 +44,11 @@ class Tool:
44
self.log = self.get_log_object()
45
if self.args and isinstance(self.args, dict):
46
for key, value in self.args.items():
47
+ ctx = {"content": str(value) if not isinstance(value, str) else value}
48
+ await call_extensions_async("tool_output_update", self.agent, ctx=ctx)
49
+ display_value = ctx["content"]
50
PrintStyle(font_color="#85C1E9", bold=True).stream(self.nice_key(key)+": ")
45
- PrintStyle(font_color="#85C1E9", padding=isinstance(value,str) and "\n" in value).stream(value)
51
+ PrintStyle(font_color="#85C1E9", padding=isinstance(value,str) and "\n" in value).stream(display_value)
52
PrintStyle().print()
53
54
async def after_execution(self, response: Response, **kwargs):