Add tests for the thread-local parent header API

Covers set_thread_parent/reset_thread_parent on OutStream,
ZMQDisplayHook, ZMQShellDisplayHook, and ZMQDisplayPublisher, and
set_thread_parent/reset_thread_parent/parent_override on
ZMQInteractiveShell:

- a thread that never set a parent falls back to the global parent
- set_thread_parent affects only the calling thread; tokens undo it
  in LIFO order
- messages (stream output, execute_result, display_data) are sent
  with the calling thread's parent header
- the shell propagates thread parents to its displayhook, display_pub,
  and sys.stdout/sys.stderr when they support it, and parent_override
  restores everything on exit, including on exceptions

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UvKd3pX2HBzbEE1y4P53gF

Author: claude

tests/test_displayhook.py

@@ -12,7 +12,7 @@
 from jupyter_client.session import Session
 from traitlets import Int
 
-from ipykernel.displayhook import ZMQShellDisplayHook
+from ipykernel.displayhook import ZMQDisplayHook, ZMQShellDisplayHook
 
 
 class NoReturnHook:
@@ -202,6 +202,107 @@ def test_empty_data_skips_send_and_hooks(self):
         assert self.session.send_count == 0
         assert self.disp.msg is None
 
+    def test_thread_parent(self):
+        """set_thread_parent affects only the calling thread's messages."""
+        parent1 = {"header": {"msg_id": "parent-1"}}
+        parent2 = {"header": {"msg_id": "parent-2"}}
+
+        self.disp.set_parent(parent1)
+        assert self.disp.parent_header == parent1["header"]
+
+        results = {}
+
+        def worker():
+            try:
+                # A thread that never set a parent falls back to the global one.
+                results["inherited"] = self.disp.parent_header
+                token = self.disp.set_thread_parent(parent2)
+                _drive(self.disp)
+                results["msg_parent"] = self.session.last_msg["parent_header"]
+                self.disp.reset_thread_parent(token)
+                results["after"] = self.disp.parent_header
+            except BaseException as e:
+                results["error"] = e
+
+        t = Thread(target=worker)
+        t.start()
+        t.join(timeout=10)
+        assert not t.is_alive()
+        assert results.get("error") is None
+
+        assert results["inherited"] == parent1["header"]
+        assert results["msg_parent"] == parent2["header"]
+        assert results["after"] == parent1["header"]
+
+        # The main thread still sends with the global parent.
+        _drive(self.disp)
+        assert self.session.last_msg["parent_header"] == parent1["header"]
+        assert self.disp.parent_header == parent1["header"]
+
+
+class ParentRecordingSession(CounterSession):
+    """Records the parent= keyword argument passed to send()."""
+
+    last_parent = None
+
+    def send(self, *args, **kwargs):
+        self.last_parent = kwargs.get("parent")
+        return super().send(*args, **kwargs)
+
+
+class ZMQDisplayHookTests(unittest.TestCase):
+    """Tests for the simple (non-shell) ZMQDisplayHook."""
+
+    def setUp(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.PUB)
+        self.session = ParentRecordingSession()
+        self.hook = ZMQDisplayHook(self.session, self.socket)
+
+    def tearDown(self):
+        self.socket.close()
+        self.context.term()
+
+    def test_thread_parent(self):
+        """set_thread_parent affects only the calling thread's messages."""
+        parent1 = {"header": {"msg_id": "parent-1"}}
+        parent2 = {"header": {"msg_id": "parent-2"}}
+
+        self.hook.set_parent(parent1)
+        assert self.hook.parent_header == parent1["header"]
+
+        self.hook(12)
+        assert self.session.last_parent == parent1["header"]
+
+        results = {}
+
+        def worker():
+            try:
+                # A thread that never set a parent falls back to the global one.
+                results["inherited"] = self.hook.parent_header
+                token = self.hook.set_thread_parent(parent2)
+                self.hook(34)
+                results["msg_parent"] = self.session.last_parent
+                self.hook.reset_thread_parent(token)
+                results["after"] = self.hook.parent_header
+            except BaseException as e:
+                results["error"] = e
+
+        t = Thread(target=worker)
+        t.start()
+        t.join(timeout=10)
+        assert not t.is_alive()
+        assert results.get("error") is None
+
+        assert results["inherited"] == parent1["header"]
+        assert results["msg_parent"] == parent2["header"]
+        assert results["after"] == parent1["header"]
+
+        # The main thread still sends with the global parent.
+        self.hook(56)
+        assert self.session.last_parent == parent1["header"]
+        assert self.hook.parent_header == parent1["header"]
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_io.py

@@ -117,6 +117,121 @@ def test_outstream(iopub_thread):
         assert stream.writable()
 
 
+class RecordingSession(Session):
+    """Session subclass that records the messages it sends."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.sent_msgs = []
+
+    def send(self, stream, msg_or_type, *args, **kwargs):
+        if isinstance(msg_or_type, dict):
+            self.sent_msgs.append(msg_or_type)
+        return super().send(stream, msg_or_type, *args, **kwargs)
+
+
+def _run_in_thread(target):
+    """Run target in a thread, re-raising any exception in the caller."""
+    error = []
+
+    def wrapper():
+        try:
+            target()
+        except BaseException as e:
+            error.append(e)
+
+    t = threading.Thread(target=wrapper)
+    t.start()
+    t.join(timeout=10)
+    assert not t.is_alive()
+    if error:
+        raise error[0]
+
+
+def test_outstream_thread_parent(iopub_thread):
+    """set_thread_parent only affects the calling thread and is undone by its token."""
+    session = Session()
+    stream = OutStream(session, iopub_thread, "stdout")
+    parent1 = {"header": {"msg_id": "parent-1"}}
+    parent2 = {"header": {"msg_id": "parent-2"}}
+    parent3 = {"header": {"msg_id": "parent-3"}}
+
+    stream.set_parent(parent1)
+    assert stream.parent_header == parent1["header"]
+
+    results = {}
+
+    def worker():
+        # A thread that never set a parent falls back to the global one.
+        results["inherited"] = stream.parent_header
+        token1 = stream.set_thread_parent(parent2)
+        results["local"] = stream.parent_header
+        # Thread parents nest, and tokens undo them in LIFO order.
+        token2 = stream.set_thread_parent(parent3)
+        results["nested"] = stream.parent_header
+        stream.reset_thread_parent(token2)
+        results["unnested"] = stream.parent_header
+        stream.reset_thread_parent(token1)
+        results["reset"] = stream.parent_header
+
+    _run_in_thread(worker)
+
+    assert results["inherited"] == parent1["header"]
+    assert results["local"] == parent2["header"]
+    assert results["nested"] == parent3["header"]
+    assert results["unnested"] == parent2["header"]
+    assert results["reset"] == parent1["header"]
+    # The main thread never saw the thread-local parents.
+    assert stream.parent_header == parent1["header"]
+
+
+def test_outstream_set_parent_is_global(iopub_thread):
+    """set_parent from any thread updates the fallback for other threads."""
+    session = Session()
+    stream = OutStream(session, iopub_thread, "stdout")
+    parent1 = {"header": {"msg_id": "parent-1"}}
+    parent2 = {"header": {"msg_id": "parent-2"}}
+
+    stream.set_parent(parent1)
+    _run_in_thread(lambda: stream.set_parent(parent2))
+
+    results = {}
+
+    def fresh_thread():
+        results["seen"] = stream.parent_header
+
+    _run_in_thread(fresh_thread)
+    assert results["seen"] == parent2["header"]
+
+
+def test_outstream_thread_parent_routing(iopub_thread):
+    """Output written under a thread parent is sent with that parent header."""
+    session = RecordingSession(key=b"abc")
+    stream = OutStream(session, iopub_thread, "stdout")
+    parent1 = {"header": {"msg_id": "parent-1"}}
+    parent2 = {"header": {"msg_id": "parent-2"}}
+
+    stream.set_parent(parent1)
+
+    def worker():
+        token = stream.set_thread_parent(parent2)
+        try:
+            stream.write("from-thread")
+        finally:
+            stream.reset_thread_parent(token)
+
+    stream.write("from-main")
+    _run_in_thread(worker)
+    stream.flush()
+
+    routed = {
+        msg["content"]["text"]: msg["parent_header"]["msg_id"]
+        for msg in session.sent_msgs
+        if msg["msg_type"] == "stream"
+    }
+    assert routed == {"from-main": "parent-1", "from-thread": "parent-2"}
+
+
 async def test_event_pipe_gc(iopub_thread):
     session = Session(key=b"abc")
     stream = OutStream(