fix(execute_code): forward session_id to nested tool calls via RPC

tests/tools/test_code_execution.py

@@ -63,7 +63,7 @@ def _fresh_kernel_registry():
 from tools.registry import registry
 
 
-def _mock_handle_function_call(function_name, function_args, task_id=None, user_task=None):
+def _mock_handle_function_call(function_name, function_args, task_id=None, user_task=None, **kwargs):
     """Mock dispatcher that returns canned responses for each tool."""
     if function_name == "terminal":
         cmd = function_args.get("command", "")
@@ -331,7 +331,7 @@ def call(i):
     print(f"OK {N}/{N}")
 '''
 
-        def slow_mock(function_name, function_args, task_id=None, user_task=None):
+        def slow_mock(function_name, function_args, task_id=None, user_task=None, **kwargs):
             import time as _t
             if function_name == "terminal":
                 _t.sleep(0.05)  # ensure requests overlap on the socket
@@ -922,5 +922,189 @@ def test_generated_module_sends_token(self):
         self.assertIn('"token"', src)
 
 
+class TestRpcSessionIdForwarding(unittest.TestCase):
+    """Regression tests for #51931: nested tool calls (invoked by
+    execute_code via RPC) must receive the parent's session_id so plugin
+    hooks (on_pre_tool_call / on_post_tool_call) can correlate them with
+    the originating turn."""
+
+    def test_registry_dispatch_forwards_session_id_to_execute_code(self):
+        """registry.dispatch must not drop the explicit session_id kwarg."""
+        from tools.registry import registry
+
+        captured = {}
+
+        def fake_execute_code(code, task_id=None, enabled_tools=None,
+                              reset=False, session_id=None):
+            captured["code"] = code
+            captured["task_id"] = task_id
+            captured["enabled_tools"] = enabled_tools
+            captured["session_id"] = session_id
+            return json.dumps({"status": "ok"})
+
+        with patch("tools.code_execution_tool.execute_code",
+                   side_effect=fake_execute_code):
+            raw = registry.dispatch(
+                "execute_code",
+                {"code": "print('hi')"},
+                task_id="test-task",
+                enabled_tools=["read_file"],
+                session_id="session-contract",
+            )
+
+        assert isinstance(raw, str), raw
+        self.assertEqual(json.loads(raw), {"status": "ok"})
+        self.assertEqual(captured, {
+            "code": "print('hi')",
+            "task_id": "test-task",
+            "enabled_tools": ["read_file"],
+            "session_id": "session-contract",
+        })
+
+    def test_rpc_server_loop_forwards_session_id(self):
+        """_rpc_server_loop must pass session_id to handle_function_call."""
+        from tools.code_execution_rpc import _rpc_server_loop
+
+        captured = {}
+
+        def fake_handle_function_call(tool_name, tool_args, task_id=None,
+                                      session_id=None, **kwargs):
+            captured["session_id"] = session_id
+            return json.dumps({"status": "ok"})
+
+        # Build a minimal mock socket that delivers one request then closes.
+        server_sock = MagicMock()
+        conn = MagicMock()
+        # Simulate one JSON request line then EOF (b"" ends the loop).
+        request = json.dumps({"tool": "read_file", "args": {"path": "/tmp/x"},
+                              "token": "test-rpc-token"})
+        conn.recv.side_effect = [(request + "\n").encode(), b""]
+        server_sock.accept.return_value = (conn, ("127.0.0.1", 12345))
+
+        stop_event = threading.Event()
+        with patch("model_tools.handle_function_call",
+                   side_effect=fake_handle_function_call):
+            _rpc_server_loop(
+                server_sock, "test-task", [], [0], 100,
+                frozenset({"read_file"}), stop_event, "test-rpc-token",
+                session_id="test-session-123",
+            )
+
+        self.assertEqual(captured.get("session_id"), "test-session-123",
+                         "session_id must be forwarded to handle_function_call")
+
+    def test_rpc_server_loop_defaults_session_id_empty(self):
+        """When session_id is not provided, it defaults to empty string
+        (backward compatibility — no crash)."""
+        from tools.code_execution_rpc import _rpc_server_loop
+
+        captured = {}
+
+        def fake_handle_function_call(tool_name, tool_args, task_id=None,
+                                      session_id=None, **kwargs):
+            captured["session_id"] = session_id
+            return json.dumps({"status": "ok"})
+
+        server_sock = MagicMock()
+        conn = MagicMock()
+        request = json.dumps({"tool": "read_file", "args": {"path": "/tmp/x"},
+                              "token": "test-rpc-token"})
+        conn.recv.side_effect = [(request + "\n").encode(), b""]
+        server_sock.accept.return_value = (conn, ("127.0.0.1", 12345))
+
+        stop_event = threading.Event()
+        with patch("model_tools.handle_function_call",
+                   side_effect=fake_handle_function_call):
+            _rpc_server_loop(
+                server_sock, "test-task", [], [0], 100,
+                frozenset({"read_file"}), stop_event, "test-rpc-token",
+                # session_id not passed — should default to ""
+            )
+
+        self.assertEqual(captured.get("session_id"), "",
+                         "session_id should default to empty string")
+
+    def test_rpc_poll_loop_forwards_session_id(self):
+        """_rpc_poll_loop (remote backend) must also forward session_id."""
+        from tools.code_execution_rpc import _rpc_poll_loop
+
+        captured = {}
+
+        def fake_handle_function_call(tool_name, tool_args, task_id=None,
+                                      session_id=None, **kwargs):
+            captured["session_id"] = session_id
+            return json.dumps({"status": "ok"})
+
+        # Build a mock env that returns one request file, then the request
+        # content, then marks it done.
+        env = MagicMock()
+        request = json.dumps({"tool": "read_file", "args": {"path": "/tmp/x"},
+                              "token": "test-rpc-token"})
+
+        # First ls: finds one request file. Subsequent ls: empty (no more).
+        ls_call_count = [0]
+
+        def env_execute(cmd, **kwargs):
+            if cmd.startswith("ls "):
+                ls_call_count[0] += 1
+                if ls_call_count[0] == 1:
+                    return {"output": "/rpc/req_001\n"}
+                return {"output": ""}
+            if cmd.startswith("cat "):
+                return {"output": request}
+            if cmd.startswith("rm "):
+                return {"output": ""}
+            return {"output": ""}
+
+        env.execute.side_effect = env_execute
+
+        stop_event = threading.Event()
+
+        def fake_handle_and_stop(*args, **kwargs):
+            result = fake_handle_function_call(*args, **kwargs)
+            # Stop the poll loop after the first dispatch.
+            stop_event.set()
+            return result
+
+        with patch("model_tools.handle_function_call",
+                   side_effect=fake_handle_and_stop):
+            _rpc_poll_loop(
+                env, "/rpc", "test-task", [], [0], 100,
+                frozenset({"read_file"}), stop_event, "test-rpc-token",
+                session_id="remote-session-456",
+            )
+
+        self.assertEqual(captured.get("session_id"), "remote-session-456",
+                         "session_id must be forwarded in remote RPC path too")
+
+    def test_execute_code_forwards_session_id_to_session_kernel(self):
+        """Local execute_code is session-kernel-only; session_id must reach it."""
+        from tools.code_execution_tool import execute_code
+
+        captured = {}
+
+        def fake_kernel(code, **kwargs):
+            captured["session_id"] = kwargs.get("session_id")
+            return json.dumps({"status": "ok"})
+
+        with patch("tools.code_kernel.execute_in_session_kernel",
+                   side_effect=fake_kernel), \
+             patch("tools.approval.check_execute_code_guard",
+                   return_value={"approved": True}), \
+             patch("tools.terminal_tool._get_env_config",
+                   return_value={"env_type": "local"}), \
+             patch("tools.process_registry._is_supervised_gateway_process",
+                   return_value=False):
+            raw = execute_code(
+                "print('hi')",
+                task_id="test-task",
+                enabled_tools=["read_file"],
+                session_id="kernel-session",
+            )
+
+        self.assertEqual(json.loads(raw), {"status": "ok"})
+        self.assertEqual(captured.get("session_id"), "kernel-session")
+
+
 if __name__ == "__main__":
     unittest.main()

tests/tools/test_code_execution_modes.py

@@ -71,7 +71,7 @@ def _mock_mode(mode):
         yield
 
 
-def _mock_handle_function_call(function_name, function_args, task_id=None, user_task=None):
+def _mock_handle_function_call(function_name, function_args, task_id=None, user_task=None, **kwargs):
     """Minimal mock dispatcher reused across tests."""
     if function_name == "terminal":
         return json.dumps({"output": "mock", "exit_code": 0})

tests/tools/test_code_kernel.py

@@ -353,18 +353,19 @@ def _cell():
                 t.join()
         self.assertEqual([r["status"] for r in results], ["success"] * 6)
         self.assertEqual(len(_KERNELS), 1)
+        # BSD pgrep (macOS) has no -c; count PID lines instead.
         live = subprocess.run(
-            ["pgrep", "-fc", "-P", str(os.getpid()), "hermes_kernel_runner"],
+            ["pgrep", "-f", "-P", str(os.getpid()), "hermes_kernel_runner"],
             capture_output=True, text=True,
-        ).stdout.strip()
-        self.assertEqual(live, "1")
+        ).stdout.split()
+        self.assertEqual(len(live), 1)
 
 
 class TestPerCellRpcAuthority(unittest.TestCase):
     """Interpreter state persists across cells; RPC authority must not."""
 
     def _recorder(self, seen):
-        def _handle(tool_name, tool_args, task_id=None):
+        def _handle(tool_name, tool_args, task_id=None, **kwargs):
             from tools.thread_context import _callback_api
 
             (get_approval, _set_a), *_rest = _callback_api()
@@ -445,6 +446,29 @@ def test_a_settled_cells_authority_refuses_dispatch(self):
         result = authority.dispatch("web_search", {"query": "q"})
         self.assertIn("No active execute_code cell", result)
 
+    def test_cell_authority_forwards_session_id(self):
+        """Nested kernel-cell tool calls must keep the parent session_id (#51931)."""
+        from tools.code_kernel import CellAuthority
+
+        captured = {}
+
+        def fake_handle(tool_name, tool_args, task_id=None, session_id=None, **kwargs):
+            captured["tool_name"] = tool_name
+            captured["task_id"] = task_id
+            captured["session_id"] = session_id
+            return json.dumps({"status": "ok"})
+
+        authority = CellAuthority("turn-1", session_id="kernel-session")
+        with patch("model_tools.handle_function_call", side_effect=fake_handle):
+            result = authority.dispatch("read_file", {"path": "/tmp/x"})
+
+        self.assertEqual(json.loads(result), {"status": "ok"})
+        self.assertEqual(captured, {
+            "tool_name": "read_file",
+            "task_id": "turn-1",
+            "session_id": "kernel-session",
+        })
+
     def test_each_cell_installs_a_fresh_authority(self):
         with _kernel_config():
             _run("x = 1")

tools/code_execution_rpc.py

@@ -25,9 +25,10 @@
 _TERMINAL_BLOCKED_PARAMS = {"background", "pty", "notify", "notify_on_complete", "watch_patterns"}
 
 
-def _default_dispatch(task_id):
+def _default_dispatch(task_id, session_id=""):
     from model_tools import handle_function_call
-    return lambda tool_name, tool_args: handle_function_call(tool_name, tool_args, task_id=task_id)
+    return lambda tool_name, tool_args: handle_function_call(
+        tool_name, tool_args, task_id=task_id, session_id=session_id)
 
 
 def _rpc_token_ok(request: dict, rpc_token: str) -> bool:
@@ -69,14 +70,22 @@ def _handle_rpc_request(request: dict, *, allowed_tools: frozenset, tool_call_co
 
 def _rpc_server_loop(server_sock: socket.socket, task_id: str, tool_call_log: list,
                      tool_call_counter: list, max_tool_calls: int, allowed_tools: frozenset,
-                     stop_event: threading.Event, rpc_token: str, dispatch=None):
+                     stop_event: threading.Event, rpc_token: str, dispatch=None,
+                     session_id: str = ""):
     """Accept one client and serve newline-delimited JSON requests until it disconnects, idles
     300s, or the call limit is reached. ``tool_call_counter`` is a mutable ``[int]``. ``dispatch``
     overrides how an allowed, budgeted call runs: per-call sandboxes use the default (the thread
     carries the cell's context); session kernels rebind each call to the CURRENT cell's authority.
+
+    ``session_id`` is forwarded to ``handle_function_call`` so nested tool calls (e.g.
+    ``read_file`` invoked by ``execute_code``) receive the same session context as the parent —
+    without it, plugin hooks ``on_pre_tool_call`` / ``on_post_tool_call`` see an empty session_id
+    and cannot correlate the nested call with the originating turn (#51931). Session kernels
+    capture this per cell on ``CellAuthority`` rather than freezing it on the long-lived serving
+    thread.
     """
     if dispatch is None:
-        dispatch = _default_dispatch(task_id)
+        dispatch = _default_dispatch(task_id, session_id=session_id)
     conn = None
     try:
         server_sock.settimeout(0.05)
@@ -129,11 +138,15 @@ def _rpc_server_loop(server_sock: socket.socket, task_id: str, tool_call_log: li
 
 def _rpc_poll_loop(env, rpc_dir: str, task_id: str, tool_call_log: list, tool_call_counter: list,
                    max_tool_calls: int, allowed_tools: frozenset, stop_event: threading.Event,
-                   rpc_token: str):
+                   rpc_token: str, session_id: str = ""):
     """Poll the remote filesystem for request files and answer them. Background thread; each
     ``env.execute()`` is an independent process, so this is safe alongside the script-execution
-    thread. Malformed or unauthorized requests are removed without a response."""
-    dispatch = _default_dispatch(task_id)
+    thread. Malformed or unauthorized requests are removed without a response.
+
+    ``session_id`` is forwarded to ``handle_function_call`` so nested tool calls receive the
+    same session context as the parent (#51931).
+    """
+    dispatch = _default_dispatch(task_id, session_id=session_id)
     poll_interval = 0.1
     quoted_rpc_dir = shlex.quote(rpc_dir)
     while not stop_event.is_set():