fix: improve buffer handling for multimodal input (#128)

* fix: improve buffer handling

* tests:add tests for StdioRequestReader

- Introduced a new private method `_read_async` to handle asynchronous reading from stdin.
- Updated `_read_stream` to utilize the new `_read_async` method for improved code clarity and maintainability.
- Added a new test file `test_stdio.py` to validate the functionality of the StdioRequestReader, ensuring correct handling of input data streams.

* fix: remove useless code

* fix: handle empty lines in StdioRequestReader

- Updated the StdioRequestReader to strip and skip empty lines during input processing.
- Added a new test to validate the handling of empty lines, ensuring the reader correctly processes a stream with mixed content.

* apply ruff

---------

Co-authored-by: Yeuoly <admin@srmxy.cn>

Author: kurokobo

python/dify_plugin/core/server/stdio/request_reader.py

@@ -17,26 +17,34 @@ class StdioRequestReader(RequestReader):
     def __init__(self):
         super().__init__()
 
+    def _read_async(self) -> bytes:
+        # read data from stdin using tp_read in 64KB chunks.
+        # the OS buffer for stdin is usually 64KB, so using a larger value doesn't make sense.
+        return tp_read(sys.stdin.fileno(), 65536)
+
     def _read_stream(self) -> Generator[PluginInStream, None, None]:
         buffer = b""
         while True:
-            # read data from stdin through tp_read
-            data = tp_read(sys.stdin.fileno(), 512)
-
+            data = self._read_async()
             if not data:
                 continue
 
             buffer += data
 
-            # process line by line and keep the last line if it is not complete
-            lines = buffer.split(b"\n")
-            if len(lines) == 0:
+            # if no b"\n" is in data, skip to the next iteration
+            if data.find(b"\n") == -1:
                 continue
 
+            # process line by line and keep the last line if it is not complete
+            lines = buffer.split(b"\n")
             buffer = lines[-1]
 
             lines = lines[:-1]
             for line in lines:
+                line = line.strip()
+                if not line:
+                    continue
+
                 try:
                     data = TypeAdapter(dict[str, Any]).validate_json(line)
                     yield PluginInStream(

python/tests/servers/test_stdio.py

@@ -0,0 +1,76 @@
+import json
+
+from dify_plugin.core.entities.plugin.io import PluginInStreamEvent
+from dify_plugin.core.server.stdio.request_reader import StdioRequestReader
+
+
+def test_stdio(monkeypatch):
+    payload = {
+        "session_id": "1",
+        "conversation_id": "2",
+        "message_id": "3",
+        "app_id": "4",
+        "endpoint_id": "5",
+        "data": {"test": "test" * 1000},
+        "event": PluginInStreamEvent.Request.value,
+    }
+
+    reader = StdioRequestReader()
+    dataflow_bytes = b"".join([json.dumps(payload).encode("utf-8") + b"\n" for _ in range(200)])
+    # split dataflow_bytes into 64KB chunks
+    dataflow_chunks = [dataflow_bytes[i : i + 65536] for i in range(0, len(dataflow_bytes), 65536)]
+
+    def mock_read_async():
+        return dataflow_chunks.pop(0)
+
+    # mock reader._read_async
+    monkeypatch.setattr(reader, "_read_async", mock_read_async)
+
+    iters = 0
+
+    for line in reader._read_stream():
+        assert line.event == PluginInStreamEvent.Request
+        assert line.session_id == "1"
+        assert line.conversation_id == "2"
+        assert line.message_id == "3"
+        assert line.app_id == "4"
+        assert line.endpoint_id == "5"
+        iters += 1
+        if iters == 200:
+            break
+
+    assert iters == 200
+
+
+def test_stdio_with_empty_line(monkeypatch):
+    payload = {
+        "session_id": "1",
+        "conversation_id": "2",
+        "message_id": "3",
+        "app_id": "4",
+        "endpoint_id": "5",
+        "data": {"test": "test" * 1000},
+        "event": PluginInStreamEvent.Request.value,
+    }
+
+    reader = StdioRequestReader()
+    dataflow_bytes = b"".join([json.dumps(payload).encode("utf-8") + b"\n" for _ in range(100)])
+    dataflow_bytes += b"\n"
+    dataflow_bytes += b"".join([json.dumps(payload).encode("utf-8") + b"\n" for _ in range(100)])
+    dataflow_bytes += b"\n"
+    dataflow_bytes += b"".join([json.dumps(payload).encode("utf-8") + b"\n" for _ in range(100)])
+    dataflow_bytes += b"\n"
+
+    def mock_read_async():
+        return dataflow_bytes
+
+    monkeypatch.setattr(reader, "_read_async", mock_read_async)
+
+    iters = 0
+    for line in reader._read_stream():
+        assert line.event == PluginInStreamEvent.Request
+        iters += 1
+        if iters == 300:
+            break
+
+    assert iters == 300