From 57f3ced290c9257c70bc4a574ec438047c1bb26e Mon Sep 17 00:00:00 2001
From: taixiansen2 <taixiansen2@gmail.com>
Date: Sun, 7 Jun 2026 13:15:44 +0800
Subject: [PATCH] Fix sensory buffer loop for oversized messages

---
 .../factory/memory_buffer/sensory_memory.py   | 34 +++++++++++--
 tests/test_sensory_memory.py                  | 49 +++++++++++++++++++
 2 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_sensory_memory.py

diff --git a/src/lightmem/factory/memory_buffer/sensory_memory.py b/src/lightmem/factory/memory_buffer/sensory_memory.py
index 8c79566..c030f16 100644
--- a/src/lightmem/factory/memory_buffer/sensory_memory.py
+++ b/src/lightmem/factory/memory_buffer/sensory_memory.py
@@ -14,17 +14,28 @@ def _recount_tokens(self) -> None:
 
     def add_messages(self, messages: List[Dict], segmenter, text_embedder) -> None:
         all_segments = []
+        if not messages:
+            return all_segments
+
         self.big_buffer.extend(messages)
 
         while self.big_buffer:
             processed_messages = []
-            for msg in self.big_buffer:
+            for pos, msg in enumerate(self.big_buffer):
                 if msg["role"] == "user":
                     cur_token_count = len(self.tokenizer.encode(msg["content"]))
                     if self.token_count + cur_token_count <= self.max_tokens:
                         self.buffer.append(msg)
                         self.token_count += cur_token_count
                         processed_messages.append(msg)
+                    elif self.token_count == 0 and not any(m.get("role") == "user" for m in self.buffer):
+                        oversize_segment = [msg]
+                        processed_messages.append(msg)
+                        if pos + 1 < len(self.big_buffer) and self.big_buffer[pos + 1].get("role") == "assistant":
+                            oversize_segment.append(self.big_buffer[pos + 1])
+                            processed_messages.append(self.big_buffer[pos + 1])
+                        all_segments.append(oversize_segment)
+                        break
                     else:
                         segments = self.cut_with_segmenter(segmenter, text_embedder)
                         all_segments.extend(segments)
@@ -47,7 +58,16 @@ def cut_with_segmenter(self, segmenter, text_embedder, force_segment: bool=False
         2. Fine-grained adjustment based on semantic similarity.
         """
         segments = []
+        if not self.buffer:
+            self.token_count = 0
+            return segments
+
         buffer_texts = [m["content"] for m in self.buffer if m["role"] == "user"]
+        if not buffer_texts:
+            self.buffer.clear()
+            self.token_count = 0
+            return segments
+
         boundaries = segmenter.propose_cut(buffer_texts)
 
         if not boundaries:
@@ -57,11 +77,19 @@ def cut_with_segmenter(self, segmenter, text_embedder, force_segment: bool=False
             return segments
 
         turns = []
-        for i in range(0, len(self.buffer), 2):
+        for i in range(0, len(self.buffer) - 1, 2):
+            if self.buffer[i].get("role") != "user" or self.buffer[i + 1].get("role") != "assistant":
+                continue
             user_msg = self.buffer[i]["content"]
             assistant_msg = self.buffer[i + 1]["content"]
             turns.append(user_msg + " " + assistant_msg)
 
+        if not turns:
+            segments.append(self.buffer.copy())
+            self.buffer.clear()
+            self.token_count = 0
+            return segments
+
         embeddings = []
         for turn in turns:
             emb = text_embedder.embed(turn)
@@ -104,7 +132,7 @@ def cut_with_segmenter(self, segmenter, text_embedder, force_segment: bool=False
 
         if force_segment:
             segments.append(self.buffer[start_idx:])
-            start_idx = len(boundaries)
+            start_idx = len(self.buffer)
 
         if start_idx > 0: 
             del self.buffer[:start_idx]
diff --git a/tests/test_sensory_memory.py b/tests/test_sensory_memory.py
new file mode 100644
index 0000000..f155e78
--- /dev/null
+++ b/tests/test_sensory_memory.py
@@ -0,0 +1,49 @@
+from lightmem.factory.memory_buffer.sensory_memory import SenMemBufferManager
+
+
+class FakeTokenizer:
+    def encode(self, text):
+        return text.split()
+
+
+class FakeSegmenter:
+    def propose_cut(self, buffer_texts):
+        return []
+
+
+class FakeEmbedder:
+    def embed(self, text):
+        return [1.0, 0.0]
+
+
+def test_oversized_single_user_message_is_consumed():
+    manager = SenMemBufferManager(max_tokens=3, tokenizer=FakeTokenizer())
+    messages = [
+        {"role": "user", "content": "one two three four"},
+        {"role": "assistant", "content": "ok"},
+    ]
+
+    segments = manager.add_messages(messages, FakeSegmenter(), FakeEmbedder())
+
+    assert segments == [messages]
+    assert manager.big_buffer == []
+    assert manager.buffer == []
+    assert manager.token_count == 0
+
+
+def test_force_segment_flushes_remaining_buffer():
+    manager = SenMemBufferManager(max_tokens=10, tokenizer=FakeTokenizer())
+    manager.buffer = [
+        {"role": "user", "content": "one"},
+        {"role": "assistant", "content": "two"},
+    ]
+    manager.token_count = 1
+
+    segments = manager.cut_with_segmenter(FakeSegmenter(), FakeEmbedder(), force_segment=True)
+
+    assert segments == [[
+        {"role": "user", "content": "one"},
+        {"role": "assistant", "content": "two"},
+    ]]
+    assert manager.buffer == []
+    assert manager.token_count == 0