fix(ai): verify OS port availability in assign_port to prevent cross-engine data-plane crosstalk (#1558)

packages/ai/src/ai/modules/task/task_server.py

@@ -67,6 +67,7 @@
 
 import time
 import asyncio
+import socket
 import uuid
 from typing import List
 from fastapi import WebSocket
@@ -695,22 +696,60 @@ def get_task(self, token: str) -> Task:
         # Extract and return the task instance
         return control.task
 
+    @staticmethod
+    def _is_port_free(port: int) -> bool:
+        """
+        Check whether a TCP port is actually free on the host.
+
+        The in-memory ``_allocated_ports`` list only knows about ports handed
+        out by *this* engine process. It cannot see ports held by another
+        engine (e.g. a second workspace, or an IDE engine alongside a CLI
+        engine) or any unrelated process. Two engines both defaulting to
+        ``base_port=20000`` would otherwise both hand out 20000/20001, and the
+        second engine's node host would silently dial the first engine's data
+        plane (see #1558).
+
+        Attempt a real bind on ``127.0.0.1:port`` to confirm availability.
+
+        Args:
+            port: TCP port number to probe.
+
+        Returns:
+            True if the port could be bound (free), False if it is in use.
+        """
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            try:
+                # No SO_REUSEADDR: we want the bind to fail if another
+                # process (or engine) is already holding this port.
+                sock.bind(('127.0.0.1', port))
+                return True
+            except OSError:
+                return False
+
     def assign_port(self) -> int:
         """
         Allocate available port from managed pool.
 
+        A candidate port must be free both in this engine's own
+        ``_allocated_ports`` list *and* on the host OS (verified with a real
+        bind), so two engines on the same machine never hand out the same
+        port and cross-connect their data planes (#1558).
+
         Returns:
             Available port number (base_port to base_port+9999 range)
 
         Raises:
             RuntimeError: If no ports available
         """
         base_port = self._config.get('base_port', 20000)
-        # Search for available port
+        # Search for a port free both in our pool and on the host
         for port in range(base_port, base_port + 10000):
-            if port not in self._allocated_ports:
-                self._allocated_ports.append(port)
-                return port
+            if port in self._allocated_ports:
+                continue
+            if not self._is_port_free(port):
+                continue
+            self._allocated_ports.append(port)
+            return port
 
         raise RuntimeError(f'No available ports in the range {base_port}-{base_port + 9999}')
 

packages/ai/tests/ai/modules/task/test_task_server.py

@@ -20,6 +20,7 @@
 
 from __future__ import annotations
 
+import socket
 import sys
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
@@ -54,6 +55,10 @@ def _make_server(*, config=None, web_server=None):
     ts._config = config if config is not None else {}
     ts._server = web_server or MagicMock()
     ts.debug_message = MagicMock()
+    # Stub the host bind-check so port-allocation *logic* tests are
+    # deterministic and never touch real OS ports. Tests that exercise the
+    # real bind behaviour override or bypass this.
+    ts._is_port_free = lambda port: True
     return ts
 
 
@@ -181,6 +186,53 @@ def test_assign_port_exhausts_after_10000():
         ts.assign_port()
 
 
+def test_assign_port_skips_os_busy_port():
+    """
+    A port free in `_allocated_ports` but busy on the host is skipped (#1558).
+
+    Simulates another engine/process already holding 20000: even though this
+    engine has never allocated it, the OS bind-check must reject it so the two
+    engines cannot hand out the same port and cross-connect their data planes.
+    """
+    ts = _make_server()
+    # 20000 is busy on the host, everything else is free.
+    ts._is_port_free = lambda port: port != 20000
+    assert ts.assign_port() == 20001
+    assert ts._allocated_ports == [20001]
+
+
+def test_is_port_free_detects_bound_port():
+    """The real bind-check reports a bound port as busy and a free one as free."""
+    # Bind a real socket to an ephemeral port and confirm the check sees it busy.
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as held:
+        held.bind(('127.0.0.1', 0))
+        held.listen(1)
+        busy_port = held.getsockname()[1]
+        assert TaskServer._is_port_free(busy_port) is False
+
+    # Once released, the same port is reported free again.
+    assert TaskServer._is_port_free(busy_port) is True
+
+
+def test_assign_port_real_bind_skips_held_port():
+    """
+    End-to-end: assign_port (using the real bind-check) skips a port held by
+    another socket and returns the next free one (#1558).
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as held:
+        held.bind(('127.0.0.1', 0))
+        held.listen(1)
+        busy_port = held.getsockname()[1]
+
+        ts = _make_server(config={'base_port': busy_port})
+        # Use the real host bind-check, not the deterministic stub.
+        del ts._is_port_free
+        assigned = ts.assign_port()
+
+    assert assigned != busy_port
+    assert assigned == busy_port + 1 or assigned > busy_port
+
+
 # ---------------------------------------------------------------------------
 # get_task_control / get_task_control_by_public_key / get_task
 # ---------------------------------------------------------------------------