From a67154c0b540b8e851f57e90b4680dae4a89c006 Mon Sep 17 00:00:00 2001 From: meetp06 <65815199+meetp06@users.noreply.github.com> Date: Mon, 13 Jul 2026 15:47:00 -0700 Subject: [PATCH] fix(ai): verify OS port availability in assign_port (#1558) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TaskServer.assign_port() allocated data/debug ports by consulting only its own in-memory `_allocated_ports` list, never checking whether the port was actually free on the host. When two engines run on one machine (two workspaces, or an IDE engine plus a CLI engine), both default to `base_port=20000` and both hand out 20000/20001 first. The second engine's node host then dials `ws://localhost:{data_port}/task/data` and connects to the *first* engine's data server, so node hosts silently talk to the wrong engine's data plane. It surfaces as a misleading downstream error (e.g. a vector-store/DNS failure), not a port collision. Add a real bind test: `_is_port_free()` attempts to bind `127.0.0.1:port` (no SO_REUSEADDR, so an in-use port fails), and `assign_port()` now skips a candidate unless it is free both in `_allocated_ports` and on the host. Two engines on the same machine can no longer hand out the same port. Tests: stub the host check in the existing logic tests so they stay deterministic, and add coverage for the new behaviour — an OS-busy port is skipped, the real bind-check reports bound/free correctly, and end-to-end assign_port skips a port held by a live socket. Co-Authored-By: Claude Opus 4.8 --- .../ai/src/ai/modules/task/task_server.py | 47 +++++++++++++++-- .../tests/ai/modules/task/test_task_server.py | 52 +++++++++++++++++++ 2 files changed, 95 insertions(+), 4 deletions(-) diff --git a/packages/ai/src/ai/modules/task/task_server.py b/packages/ai/src/ai/modules/task/task_server.py index 950a4ef8e..c6e5b0d85 100644 --- a/packages/ai/src/ai/modules/task/task_server.py +++ b/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,10 +696,45 @@ 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) @@ -706,11 +742,14 @@ def assign_port(self) -> int: 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}') diff --git a/packages/ai/tests/ai/modules/task/test_task_server.py b/packages/ai/tests/ai/modules/task/test_task_server.py index 07d77e3da..55b896792 100644 --- a/packages/ai/tests/ai/modules/task/test_task_server.py +++ b/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 # ---------------------------------------------------------------------------