Merge pull request #184 from openqj/feat-mcp-stdin-starve-main

fix: MCP 主循环 stdin 批量行饥饿

cloud/playwright_mcp.py

@@ -372,8 +372,56 @@ def _drain_remote(browser):
               file=sys.stderr)
 
 
+class _LineReader:
+    """Non-blocking stdin line pump (starvation-free).
+
+    The previous loop used select() + sys.stdin.readline(). That
+    starves on pipelined input: one readline() can buffer several
+    lines in userspace, after which select() sees an empty fd
+    forever and the buffered lines are never read. A client that
+    writes notifications/initialized and its next request
+    back-to-back (every MCP client does this after initialize)
+    hangs the server exactly this way. Reading the raw fd into
+    our own buffer delivers every line in arrival order.
+    """
+
+    def __init__(self, fd):
+        self._fd = fd
+        os.set_blocking(fd, False)
+        self._buf = b""
+        self.eof = False
+
+    def poll(self):
+        """Return (lines, drained).
+
+        lines are decoded strings without the newline. A trailing
+        partial line is returned once at EOF. drained is True only
+        when EOF was reached and nothing is left buffered.
+        """
+        if not self.eof:
+            try:
+                chunk = os.read(self._fd, 65536)
+            except BlockingIOError:
+                chunk = None
+            except OSError:
+                self.eof = True
+                chunk = None
+            if chunk == b"":
+                self.eof = True
+            elif chunk:
+                self._buf += chunk
+        lines = []
+        while b"\n" in self._buf:
+            raw, self._buf = self._buf.split(b"\n", 1)
+            lines.append(raw.decode("utf-8", "replace"))
+        if self.eof and self._buf:
+            lines.append(self._buf.decode("utf-8", "replace"))
+            self._buf = b""
+        return lines, (self.eof and not self._buf)
+
+
 def main():
-    import select
+    import time
     browser = _get_browser()
     # Frame sink must be set before any navigate so pump_screencast
     # has frames for the bridge to publish. No bridge (missing
@@ -381,42 +429,44 @@ def main():
     # of accumulating unboundedly in a queue nobody drains.
     if _get_bridge(browser) is not None:
         browser.set_frame_sink(_frame_queue())
-    buffer = ""
+    reader = _LineReader(sys.stdin.fileno())
     while True:
         _drain_remote(browser)
-        try:
-            readable, _, _ = select.select([sys.stdin], [], [], 0.05)
-        except Exception:  # noqa: BLE001 - stdin gone
+        lines, drained = reader.poll()
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+            _handle_line(line, browser)
+        if drained:
             break
-        if not readable:
-            continue
-        chunk = sys.stdin.readline()
-        if not chunk:
-            break
-        line = chunk.strip()
-        if not line:
-            continue
-        try:
-            request = json.loads(line)
-        except json.JSONDecodeError as exc:
-            sys.stdout.write(json.dumps(
-                {"jsonrpc": "2.0", "id": None,
-                 "error": _err(-32700, "parse error: %s" % exc)}) + "\n")
-            sys.stdout.flush()
-            continue
-        try:
-            response = handle_request(request, browser)
-        except Exception as exc:  # noqa: BLE001 - server must not die
-            print("playwright_mcp: handler error: %s" % exc,
-                  file=sys.stderr)
-            if "id" in request:
-                response = {"jsonrpc": "2.0", "id": request["id"],
-                            "error": _err(-32603, "internal error")}
-            else:
-                response = None
-        if response is not None:
-            sys.stdout.write(json.dumps(response) + "\n")
-            sys.stdout.flush()
+        if not lines:
+            time.sleep(0.05)
+
+
+def _handle_line(line, browser):
+    """Handle one JSON-RPC line; responses go to stdout."""
+    try:
+        request = json.loads(line)
+    except json.JSONDecodeError as exc:
+        sys.stdout.write(json.dumps(
+            {"jsonrpc": "2.0", "id": None,
+             "error": _err(-32700, "parse error: %s" % exc)}) + "\n")
+        sys.stdout.flush()
+        return
+    try:
+        response = handle_request(request, browser)
+    except Exception as exc:  # noqa: BLE001 - server must not die
+        print("playwright_mcp: handler error: %s" % exc,
+              file=sys.stderr)
+        if "id" in request:
+            response = {"jsonrpc": "2.0", "id": request["id"],
+                        "error": _err(-32603, "internal error")}
+        else:
+            response = None
+    if response is not None:
+        sys.stdout.write(json.dumps(response) + "\n")
+        sys.stdout.flush()
 
 
 def _frame_queue():

cloud/tests/test_playwright.py

@@ -359,5 +359,63 @@ class MCPTest(unittest.TestCase):
             self.assertIn(name, names)
 
 
+class LineReaderTest(unittest.TestCase):
+    """Pipelined stdin lines must all arrive (no select starvation).
+
+    Regression test for the live failure where the MCP answered
+    initialize and then hung forever: the client wrote
+    notifications/initialized + tools/call back-to-back, one
+    readline() buffered both, and select() never saw the second
+    line again because the fd was drained.
+    """
+
+    def test_pipelined_lines_all_delivered(self):
+        import time
+        from playwright_mcp import _LineReader
+        rfd, wfd = os.pipe()
+        try:
+            reader = _LineReader(rfd)
+            os.write(wfd, b"init\n")
+            lines, drained = reader.poll()
+            self.assertEqual(lines, ["init"])
+            self.assertFalse(drained)
+            # Second batch arrives as one write, like a real client
+            # sending notification + next request back-to-back.
+            os.write(wfd, b"notif\nnavigate\n")
+            # Simulate the drain-loop delay: both lines sit in the
+            # pipe before the next poll.
+            time.sleep(0.2)
+            got = []
+            deadline = time.time() + 2
+            while len(got) < 2 and time.time() < deadline:
+                lines, drained = reader.poll()
+                got.extend(lines)
+                if not lines:
+                    time.sleep(0.02)
+            self.assertEqual(got, ["notif", "navigate"])
+        finally:
+            os.close(wfd)
+            os.close(rfd)
+
+    def test_partial_line_flushed_at_eof(self):
+        from playwright_mcp import _LineReader
+        rfd, wfd = os.pipe()
+        try:
+            reader = _LineReader(rfd)
+            os.write(wfd, b"partial-no-newline")
+            os.close(wfd)
+            lines, drained = reader.poll()
+            self.assertEqual(lines, [])
+            self.assertFalse(drained)
+            lines, drained = reader.poll()
+            self.assertEqual(lines, ["partial-no-newline"])
+            self.assertTrue(drained)
+        finally:
+            try:
+                os.close(rfd)
+            except OSError:
+                pass
+
+
 if __name__ == "__main__":
     unittest.main()