Install session start hooks for automatic agent-toolkit skill updates

awscli/customizations/agenttoolkit/atomic.py

@@ -116,18 +116,21 @@ def _rename(src, dst):
     _call_with_windows_retry(os.rename, src, dst)
 
 
-def write_file_atomic(path, content, prefix):
+def write_file_atomic(path, content, prefix, mode=None):
     """Replace ``path`` with ``content`` without ever leaving it partial.
 
     The temporary file is created as ``<prefix>*.tmp`` next to ``path``
-    with owner-only permissions, and removed if anything fails.
+    with owner-only permissions, and removed if anything fails. ``content``
+    is written as is, without newline translation.
     """
     fd, tmp_path = tempfile.mkstemp(
         dir=os.path.dirname(path), prefix=prefix, suffix='.tmp'
     )
     try:
-        with os.fdopen(fd, 'w', encoding='utf-8') as f:
+        with os.fdopen(fd, 'w', encoding='utf-8', newline='') as f:
             f.write(content)
+        if mode is not None:
+            os.chmod(tmp_path, mode)
         _call_with_windows_retry(os.replace, tmp_path, path)
     except BaseException:
         with contextlib.suppress(OSError):

awscli/customizations/agenttoolkit/configure.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 import io
+import os
 import sys
 
 from awscli.customizations.agenttoolkit.agents import (
@@ -20,6 +21,22 @@
     collapse_home,
     universal_first,
 )
+from awscli.customizations.agenttoolkit.hooks import (
+    CODEX_TRUST_NOTE,
+    KIRO_AGENT_NOTE,
+    KIRO_NO_AGENT_NOTE,
+    KIRO_VERSION_NOTE,
+    HookAction,
+    format_kiro_agent_manual_entry,
+    format_manual_entry,
+    hook_config_path,
+    install_hook,
+    install_kiro_agent_hook,
+    is_hook_supported,
+    kiro_default_agent_path,
+    remove_hook,
+    remove_kiro_agent_hooks,
+)
 from awscli.customizations.agenttoolkit.state import (
     StateFileError,
     is_auto_update_enabled,
@@ -48,6 +65,12 @@
 )
 
 
+_INSTALL_STATUS = {
+    HookAction.INSTALLED: 'hook installed',
+    HookAction.ALREADY_INSTALLED: 'already installed',
+}
+
+
 class ConfigureAgentToolkitCommand(BasicCommand):
     NAME = 'agent-toolkit'
     DESCRIPTION = (
@@ -154,8 +177,12 @@ def _run_main(self, parsed_args, parsed_globals):
         client = self._client or create_client(self._session, parsed_globals)
         self._install_default_skills(install_targets, client, yes=yes)
         self._configure_mcp(selected_agents, yes=yes)
-        if not no_auto_update:
-            self._configure_auto_update(install_targets, yes=yes)
+        if no_auto_update:
+            self._remove_hooks(selected_agents)
+        else:
+            self._configure_auto_update(
+                selected_agents, install_targets, yes=yes
+            )
 
         uni_print(
             '\nYou can discover additional skills with '
@@ -224,12 +251,12 @@ def _install_default_skills(self, selected_agents, client, yes=False):
                 '\nNo skills were installed successfully.\n', self._stream
             )
 
-    def _configure_auto_update(self, agents, yes=False):
+    def _configure_auto_update(self, agents, skill_targets, yes=False):
         # This wizard is the only place allowed to change the setting, so it
         # asks whenever it is on, letting customers opt out as well as in.
         currently_enabled = is_auto_update_enabled()
         if not currently_enabled and not any(
-            agent.get_installed_skills() for agent in agents
+            agent.get_installed_skills() for agent in skill_targets
         ):
             # Nothing to keep up to date
             return
@@ -242,23 +269,168 @@ def _configure_auto_update(self, agents, yes=False):
                 else AUTO_UPDATE_PROMPT
             )
             enabled = yes_no_choice(prompt, default=True)
-        self._save_auto_update(enabled)
+        if not self._save_auto_update(enabled):
+            return
+        if enabled:
+            self._install_hooks(agents)
+        else:
+            self._remove_hooks(agents)
 
     def _save_auto_update(self, enabled):
         if enabled == is_auto_update_enabled():
-            return
+            return True
         try:
             set_auto_update(enabled)
         except StateFileError as e:
             uni_print(
                 f'\nCould not save the automatic skill update setting to {e}\n',
                 self._stream,
             )
-            return
+            return False
         if enabled:
             uni_print('\nAutomatic skill updates enabled.\n', self._stream)
         else:
             uni_print('\nAutomatic skill updates turned off.\n', self._stream)
+        return True
+
+    def _install_hooks(self, agents):
+        uni_print('\nAutomatic update hooks:\n', self._stream)
+        unsupported = []
+        for agent in agents:
+            if not is_hook_supported(agent.config.id):
+                unsupported.append(agent)
+                continue
+            try:
+                action, path = install_hook(agent)
+            except OSError as e:
+                path = hook_config_path(agent.config, agent.base_dir)
+                self._print_hook_line(
+                    '\u2717', agent, path, f'could not install the hook ({e})'
+                )
+                continue
+            if action is HookAction.UNREADABLE:
+                self._print_hook_line(
+                    '\u2717',
+                    agent,
+                    path,
+                    'could not be updated because it is not plain JSON',
+                )
+                uni_print(
+                    format_manual_entry(agent.config.id, collapse_home(path)),
+                    self._stream,
+                )
+                continue
+            self._print_hook_line(
+                '\u2713', agent, path, _INSTALL_STATUS[action]
+            )
+            if agent.config.id == 'codex':
+                uni_print(f'    {CODEX_TRUST_NOTE}\n', self._stream)
+            elif agent.config.id == 'kiro':
+                uni_print(f'    {KIRO_VERSION_NOTE}\n', self._stream)
+                self._install_kiro_agent_hook(agent)
+        for agent in unsupported:
+            uni_print(
+                f'  \u2717 {agent.display_name} '
+                '(automatic updates not supported yet)\n',
+                self._stream,
+            )
+
+    def _install_kiro_agent_hook(self, agent):
+        try:
+            action, path = install_kiro_agent_hook(agent)
+        except OSError as e:
+            path = kiro_default_agent_path(agent) or os.path.join(
+                agent.base_dir, 'agents'
+            )
+            self._print_hook_line(
+                '\u2717', agent, path, f'could not install the hook ({e})'
+            )
+            return
+        if path is None:
+            uni_print(
+                f'  \u2717 {agent.display_name} \u2014 {KIRO_NO_AGENT_NOTE}\n'
+                + format_kiro_agent_manual_entry(
+                    collapse_home(
+                        os.path.join(agent.base_dir, 'agents', '<agent>.json')
+                    )
+                ),
+                self._stream,
+            )
+            return
+        if action is HookAction.UNREADABLE:
+            self._print_hook_line(
+                '\u2717',
+                agent,
+                path,
+                'could not be updated because it is not plain JSON',
+            )
+            uni_print(
+                format_kiro_agent_manual_entry(collapse_home(path)),
+                self._stream,
+            )
+        elif action is HookAction.NOT_NEEDED:
+            self._print_hook_line(
+                '\u2713', agent, path, 'not needed for V3 agents'
+            )
+            return
+        else:
+            self._print_hook_line(
+                '\u2713', agent, path, _INSTALL_STATUS[action]
+            )
+        uni_print(f'    {KIRO_AGENT_NOTE}\n', self._stream)
+
+    def _remove_hooks(self, agents):
+        lines = []
+        for agent in agents:
+            config = agent.config
+            if not is_hook_supported(config.id):
+                continue
+            try:
+                results = [remove_hook(agent)]
+            except OSError as e:
+                path = hook_config_path(agent.config, agent.base_dir)
+                lines.append(
+                    ('\u2717', agent, path, f'could not remove the hook ({e})')
+                )
+                continue
+            if config.id == 'kiro':
+                try:
+                    results.extend(remove_kiro_agent_hooks(agent))
+                except OSError as e:
+                    path = os.path.join(agent.base_dir, 'agents')
+                    lines.append(
+                        (
+                            '\u2717',
+                            agent,
+                            path,
+                            f'could not remove the hook ({e})',
+                        )
+                    )
+            for action, path in results:
+                if action is HookAction.REMOVED:
+                    lines.append(('\u2713', agent, path, 'hook removed'))
+                elif action is HookAction.UNREADABLE:
+                    lines.append(
+                        (
+                            '\u2717',
+                            agent,
+                            path,
+                            'could not be updated; remove the "aws '
+                            'agent-toolkit update-skill" hook from it by hand',
+                        )
+                    )
+        if not lines:
+            return
+        uni_print('\nAutomatic update hooks:\n', self._stream)
+        for line in lines:
+            self._print_hook_line(*line)
+
+    def _print_hook_line(self, mark, agent, path, status):
+        uni_print(
+            f'  {mark} {agent.display_name} \u2014 '
+            f'{collapse_home(path)}: {status}\n',
+            self._stream,
+        )
 
     def _configure_mcp(self, agents, yes=False):
         if not yes and not yes_no_choice(