fix: scope one-shot event handlers to registrations

src/openai/_event_handler.py

@@ -10,8 +10,7 @@ class EventHandlerRegistry:
     """Thread-safe (optional) registry of event handlers."""
 
     def __init__(self, *, use_lock: bool = False) -> None:
-        self._handlers: dict[str, list[EventHandler]] = {}
-        self._once_ids: set[int] = set()
+        self._handlers: dict[str, list[tuple[EventHandler, bool]]] = {}
         self._lock: threading.Lock | None = threading.Lock() if use_lock else None
 
     def _acquire(self) -> None:
@@ -26,9 +25,7 @@ def add(self, event_type: str, handler: EventHandler, *, once: bool = False) ->
         self._acquire()
         try:
             handlers = self._handlers.setdefault(event_type, [])
-            handlers.append(handler)
-            if once:
-                self._once_ids.add(id(handler))
+            handlers.append((handler, once))
         finally:
             self._release()
 
@@ -37,11 +34,10 @@ def remove(self, event_type: str, handler: EventHandler) -> None:
         try:
             handlers = self._handlers.get(event_type)
             if handlers is not None:
-                try:
-                    handlers.remove(handler)
-                except ValueError:
-                    pass
-                self._once_ids.discard(id(handler))
+                for index, (registered_handler, _) in enumerate(handlers):
+                    if registered_handler is handler or registered_handler == handler:
+                        del handlers[index]
+                        break
         finally:
             self._release()
 
@@ -52,11 +48,8 @@ def get_handlers(self, event_type: str) -> list[EventHandler]:
             handlers = self._handlers.get(event_type)
             if not handlers:
                 return []
-            result = list(handlers)
-            to_remove = [h for h in result if id(h) in self._once_ids]
-            for h in to_remove:
-                handlers.remove(h)
-                self._once_ids.discard(id(h))
+            result = [handler for handler, _ in handlers]
+            self._handlers[event_type] = [(handler, once) for handler, once in handlers if not once]
             return result
         finally:
             self._release()
@@ -74,10 +67,8 @@ def merge_into(self, target: EventHandlerRegistry) -> None:
         self._acquire()
         try:
             for event_type, handlers in self._handlers.items():
-                for handler in handlers:
-                    once = id(handler) in self._once_ids
+                for handler, once in handlers:
                     target.add(event_type, handler, once=once)
             self._handlers.clear()
-            self._once_ids.clear()
         finally:
             self._release()

tests/test_event_handler.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from openai._event_handler import EventHandlerRegistry
+
+
+def test_once_registration_is_scoped_to_its_event() -> None:
+    def handler(_event: object) -> None:
+        pass
+
+    registry = EventHandlerRegistry()
+    registry.add("persistent", handler)
+    registry.add("one-shot", handler, once=True)
+
+    assert registry.get_handlers("persistent") == [handler]
+    assert registry.get_handlers("persistent") == [handler]
+    assert registry.get_handlers("one-shot") == [handler]
+    assert registry.get_handlers("one-shot") == []
+
+
+def test_remove_only_removes_the_first_matching_registration() -> None:
+    def handler(_event: object) -> None:
+        pass
+
+    registry = EventHandlerRegistry()
+    registry.add("event", handler)
+    registry.add("event", handler, once=True)
+
+    registry.remove("event", handler)
+
+    assert registry.get_handlers("event") == [handler]
+    assert registry.get_handlers("event") == []
+
+
+def test_remove_prefers_identity_to_overloaded_equality() -> None:
+    class Handler:
+        def __call__(self, _event: object) -> None:
+            pass
+
+        def __eq__(self, other: object) -> bool:
+            del other
+            raise AssertionError("identity matches must not invoke __eq__")
+
+    handler = Handler()
+    registry = EventHandlerRegistry()
+    registry.add("event", handler)
+
+    registry.remove("event", handler)
+
+    assert registry.has_handlers("event") is False
+
+
+def test_merge_preserves_each_registration_mode() -> None:
+    def handler(_event: object) -> None:
+        pass
+
+    source = EventHandlerRegistry()
+    target = EventHandlerRegistry()
+    source.add("event", handler)
+    source.add("event", handler, once=True)
+
+    source.merge_into(target)
+
+    assert source.has_handlers("event") is False
+    assert target.get_handlers("event") == [handler, handler]
+    assert target.get_handlers("event") == [handler]