From 26789a07595ed50fbf210e6afa5be15d71b3e714 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E4=BA=AC=E5=9D=A4?= <2802629917@qq.com>
Date: Fri, 23 Oct 2026 20:29:56 +0800
Subject: [PATCH] fix: pause Instructor callbacks when agent hooks are disabled

---
 .../atomic_agents/agents/atomic_agent.py      | 16 ++++-
 .../tests/agents/test_agent_hooks.py          | 64 +++++++++++++++++++
 2 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py
index 76d2423b..2ac9381f 100644
--- a/atomic-agents/atomic_agents/agents/atomic_agent.py
+++ b/atomic-agents/atomic_agents/agents/atomic_agent.py
@@ -774,7 +774,7 @@ def register_hook(self, event: str, handler: Callable) -> None:
         self._hook_handlers[event].append(handler)
 
         # Register with instructor client if it supports hooks
-        if hasattr(self.client, "on"):
+        if self._hooks_enabled and hasattr(self.client, "on"):
             self.client.on(event, handler)
 
     def unregister_hook(self, event: str, handler: Callable) -> None:
@@ -789,7 +789,7 @@ def unregister_hook(self, event: str, handler: Callable) -> None:
             self._hook_handlers[event].remove(handler)
 
             # Remove from instructor client if it supports hooks
-            if hasattr(self.client, "off"):
+            if self._hooks_enabled and hasattr(self.client, "off"):
                 self.client.off(event, handler)
 
     def clear_hooks(self, event: Optional[str] = None) -> None:
@@ -828,11 +828,23 @@ def _dispatch_hook(self, event: str, *args, **kwargs) -> None:
 
     def enable_hooks(self) -> None:
         """Enable hook processing."""
+        if self._hooks_enabled:
+            return
         self._hooks_enabled = True
+        if hasattr(self.client, "on"):
+            for event, handlers in self._hook_handlers.items():
+                for handler in handlers:
+                    self.client.on(event, handler)
 
     def disable_hooks(self) -> None:
         """Disable hook processing."""
+        if not self._hooks_enabled:
+            return
         self._hooks_enabled = False
+        if hasattr(self.client, "off"):
+            for event, handlers in self._hook_handlers.items():
+                for handler in handlers:
+                    self.client.off(event, handler)
 
     @property
     def hooks_enabled(self) -> bool:
diff --git a/atomic-agents/tests/agents/test_agent_hooks.py b/atomic-agents/tests/agents/test_agent_hooks.py
index 0afedfbf..e08c3b71 100644
--- a/atomic-agents/tests/agents/test_agent_hooks.py
+++ b/atomic-agents/tests/agents/test_agent_hooks.py
@@ -24,3 +24,67 @@ def test_clear_hooks_preserves_foreign_handlers(event):
     owned.assert_not_called()
     assert other_owned.call_count == (0 if event is None else 1)
     assert agent._hook_handlers.get("completion:kwargs", []) == []
+
+
+def test_instructor_hooks_stop_and_resume_without_duplicates():
+    client = instructor.Instructor(client=Mock(), create=Mock())
+    agent = AtomicAgent(AgentConfig(client=client))
+    foreign = Mock()
+    owned = Mock()
+    client.on("completion:kwargs", foreign)
+    agent.register_hook("completion:kwargs", owned)
+
+    client.hooks.emit_completion_arguments(model="before")
+    owned.assert_called_once_with(model="before")
+
+    agent.disable_hooks()
+    agent.disable_hooks()
+    client.hooks.emit_completion_arguments(model="disabled")
+    assert owned.call_count == 1
+    assert foreign.call_count == 2
+
+    agent.enable_hooks()
+    agent.enable_hooks()
+    client.hooks.emit_completion_arguments(model="resumed")
+    assert owned.call_count == 2
+    owned.assert_called_with(model="resumed")
+    assert foreign.call_count == 3
+
+
+def test_instructor_hooks_registered_while_disabled_wait_for_enable():
+    client = instructor.Instructor(client=Mock(), create=Mock())
+    agent = AtomicAgent(AgentConfig(client=client))
+    handler = Mock()
+    agent.disable_hooks()
+    agent.register_hook("completion:kwargs", handler)
+
+    client.hooks.emit_completion_arguments(model="disabled")
+    handler.assert_not_called()
+
+    agent.enable_hooks()
+    client.hooks.emit_completion_arguments(model="resumed")
+    handler.assert_called_once_with(model="resumed")
+
+
+@pytest.mark.parametrize("remove", ["unregister", "clear_event", "clear_all"])
+def test_removing_disabled_instructor_hooks_does_not_restore_them(remove):
+    client = instructor.Instructor(client=Mock(), create=Mock())
+    agent = AtomicAgent(AgentConfig(client=client))
+    removed = Mock()
+    retained = Mock()
+    agent.register_hook("completion:kwargs", removed)
+    agent.register_hook("completion:response", retained)
+    agent.disable_hooks()
+
+    if remove == "unregister":
+        agent.unregister_hook("completion:kwargs", removed)
+    elif remove == "clear_event":
+        agent.clear_hooks("completion:kwargs")
+    else:
+        agent.clear_hooks()
+
+    agent.enable_hooks()
+    client.hooks.emit_completion_arguments(model="resumed")
+    client.hooks.emit_completion_response("response")
+    removed.assert_not_called()
+    assert retained.call_count == (0 if remove == "clear_all" else 1)