From 9bddcd474506bcbedad1fcacba31850e90aacc86 Mon Sep 17 00:00:00 2001
From: Yeuoly <admin@srmxy.cn>
Date: Fri, 15 Nov 2024 19:27:35 +0800
Subject: [PATCH] fix: decoding error when chunking messages

---
 .../dify_plugin/core/server/stdio/request_reader.py   | 11 ++++++-----
 python/dify_plugin/core/server/tcp/request_reader.py  | 10 +++++-----
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/python/dify_plugin/core/server/stdio/request_reader.py b/python/dify_plugin/core/server/stdio/request_reader.py
index 4239f03..97ad8bf 100644
--- a/python/dify_plugin/core/server/stdio/request_reader.py
+++ b/python/dify_plugin/core/server/stdio/request_reader.py
@@ -14,24 +14,25 @@
 
 class StdioRequestReader(RequestReader):
     def _read_stream(self) -> Generator[PluginInStream, None, None]:
-        buffer = ""
+        buffer = b""
         while True:
             # read data from stdin through tp_read
-            data = tp_read(sys.stdin.fileno(), 512).decode()
+            data = tp_read(sys.stdin.fileno(), 512)
 
             if not data:
                 continue
+
             buffer += data
 
             # process line by line and keep the last line if it is not complete
-            lines = buffer.split("\n")
+            lines = buffer.split(b"\n")
             if len(lines) == 0:
                 continue
 
-            if lines[-1] != "":
+            if lines[-1] != b"":
                 buffer = lines[-1]
             else:
-                buffer = ""
+                buffer = b""
 
             lines = lines[:-1]
             for line in lines:
diff --git a/python/dify_plugin/core/server/tcp/request_reader.py b/python/dify_plugin/core/server/tcp/request_reader.py
index da74eac..7ba28cd 100644
--- a/python/dify_plugin/core/server/tcp/request_reader.py
+++ b/python/dify_plugin/core/server/tcp/request_reader.py
@@ -108,14 +108,14 @@ def _read_stream(self) -> Generator[PluginInStream, None, None]:
         """
         Read data from the target
         """
-        buffer = ""
+        buffer = b""
         while self.alive:
             try:
                 ready_to_read, _, _ = select([self.sock], [], [], 1)
                 if not ready_to_read:
                     continue
-                data = self.sock.recv(4096).decode()
-                if data == "":
+                data = self.sock.recv(4096)
+                if data == b"":
                     raise Exception("Connection is closed")
             except Exception as e:
                 logger.error(f"\033[31mFailed to read data from {self.host}:{self.port}, {e}\033[0m")
@@ -130,14 +130,14 @@ def _read_stream(self) -> Generator[PluginInStream, None, None]:
             buffer += data
 
             # process line by line and keep the last line if it is not complete
-            lines = buffer.split("\n")
+            lines = buffer.split(b"\n")
             if len(lines) == 0:
                 continue
 
             if lines[-1] != "":
                 buffer = lines[-1]
             else:
-                buffer = ""
+                buffer = b""
 
             lines = lines[:-1]
             for line in lines: