Merge pull request #184 from openqj/feat-mcp-stdin-starve-main
fix: MCP 主循环 stdin 批量行饥饿
This commit is contained in:
2 files changed
+142
-34
No files matched your search
+84
-34
@@ -372,8 +372,56 @@ def _drain_remote(browser):
|
|||||||
file=sys.stderr)
|
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():
|
def main():
|
||||||
import select
|
import time
|
||||||
browser = _get_browser()
|
browser = _get_browser()
|
||||||
# Frame sink must be set before any navigate so pump_screencast
|
# Frame sink must be set before any navigate so pump_screencast
|
||||||
# has frames for the bridge to publish. No bridge (missing
|
# has frames for the bridge to publish. No bridge (missing
|
||||||
@@ -381,42 +429,44 @@ def main():
|
|||||||
# of accumulating unboundedly in a queue nobody drains.
|
# of accumulating unboundedly in a queue nobody drains.
|
||||||
if _get_bridge(browser) is not None:
|
if _get_bridge(browser) is not None:
|
||||||
browser.set_frame_sink(_frame_queue())
|
browser.set_frame_sink(_frame_queue())
|
||||||
buffer = ""
|
reader = _LineReader(sys.stdin.fileno())
|
||||||
while True:
|
while True:
|
||||||
_drain_remote(browser)
|
_drain_remote(browser)
|
||||||
try:
|
lines, drained = reader.poll()
|
||||||
readable, _, _ = select.select([sys.stdin], [], [], 0.05)
|
for line in lines:
|
||||||
except Exception: # noqa: BLE001 - stdin gone
|
line = line.strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
_handle_line(line, browser)
|
||||||
|
if drained:
|
||||||
break
|
break
|
||||||
if not readable:
|
if not lines:
|
||||||
continue
|
time.sleep(0.05)
|
||||||
chunk = sys.stdin.readline()
|
|
||||||
if not chunk:
|
|
||||||
break
|
def _handle_line(line, browser):
|
||||||
line = chunk.strip()
|
"""Handle one JSON-RPC line; responses go to stdout."""
|
||||||
if not line:
|
try:
|
||||||
continue
|
request = json.loads(line)
|
||||||
try:
|
except json.JSONDecodeError as exc:
|
||||||
request = json.loads(line)
|
sys.stdout.write(json.dumps(
|
||||||
except json.JSONDecodeError as exc:
|
{"jsonrpc": "2.0", "id": None,
|
||||||
sys.stdout.write(json.dumps(
|
"error": _err(-32700, "parse error: %s" % exc)}) + "\n")
|
||||||
{"jsonrpc": "2.0", "id": None,
|
sys.stdout.flush()
|
||||||
"error": _err(-32700, "parse error: %s" % exc)}) + "\n")
|
return
|
||||||
sys.stdout.flush()
|
try:
|
||||||
continue
|
response = handle_request(request, browser)
|
||||||
try:
|
except Exception as exc: # noqa: BLE001 - server must not die
|
||||||
response = handle_request(request, browser)
|
print("playwright_mcp: handler error: %s" % exc,
|
||||||
except Exception as exc: # noqa: BLE001 - server must not die
|
file=sys.stderr)
|
||||||
print("playwright_mcp: handler error: %s" % exc,
|
if "id" in request:
|
||||||
file=sys.stderr)
|
response = {"jsonrpc": "2.0", "id": request["id"],
|
||||||
if "id" in request:
|
"error": _err(-32603, "internal error")}
|
||||||
response = {"jsonrpc": "2.0", "id": request["id"],
|
else:
|
||||||
"error": _err(-32603, "internal error")}
|
response = None
|
||||||
else:
|
if response is not None:
|
||||||
response = None
|
sys.stdout.write(json.dumps(response) + "\n")
|
||||||
if response is not None:
|
sys.stdout.flush()
|
||||||
sys.stdout.write(json.dumps(response) + "\n")
|
|
||||||
sys.stdout.flush()
|
|
||||||
|
|
||||||
|
|
||||||
def _frame_queue():
|
def _frame_queue():
|
||||||
|
|||||||
@@ -359,5 +359,63 @@ class MCPTest(unittest.TestCase):
|
|||||||
self.assertIn(name, names)
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
Reference in new issue
Block a user