fix(compaction): remove think part from constructed compact message (#447)

Signed-off-by: Richard Chien <[email protected]>
Signed-off-by: stdrc <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: stdrc

CHANGELOG.md

@@ -17,6 +17,7 @@ Only write entries that are worth mentioning to users.
 - Core: Fix startup crash when there is broken symbolic link in the working directory
 - Core: Add builtin `okabe` agent file with `SendDMail` tool enabled
 - CLI: Add `--agent` option to specify builtin agents like `default` and `okabe`
+- Core: Improve compaction logic to better preserve relevant information
 
 ## [0.61] - 2025-12-04
 

src/kimi_cli/prompts/compact.md

@@ -1,4 +1,7 @@
-You are tasked with compacting a coding conversation context. This is critical for maintaining an effective working memory for the coding agent.
+
+---
+
+The above is a list of messages in an agent conversation. You are now given a task to compact this conversation context according to specific priorities and rules.
 
 **Compression Priorities (in order):**
 1. **Current Task State**: What is being worked on RIGHT NOW
@@ -19,10 +22,6 @@ You are tasked with compacting a coding conversation context. This is critical f
 - For errors: Keep full error message + final solution
 - For discussions: Extract decisions and action items only
 
-**Input Context to Compress:**
-
-${CONTEXT}
-
 **Required Output Structure:**
 
 <current_focus>

src/kimi_cli/soul/compaction.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from string import Template
-from typing import TYPE_CHECKING, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, NamedTuple, Protocol, runtime_checkable
 
 import kosong
-from kosong.message import ContentPart, Message, ThinkPart
+from kosong.message import ContentPart, Message, TextPart, ThinkPart
 from kosong.tooling.empty import EmptyToolset
 
 import kimi_cli.prompts as prompts
@@ -33,46 +32,21 @@ async def compact(self, messages: Sequence[Message], llm: LLM) -> Sequence[Messa
         ...
 
 
-class SimpleCompaction(Compaction):
-    MAX_PRESERVED_MESSAGES = 2
-
-    async def compact(self, messages: Sequence[Message], llm: LLM) -> Sequence[Message]:
-        history = list(messages)
-        if not history:
-            return history
+if TYPE_CHECKING:
 
-        preserve_start_index = len(history)
-        n_preserved = 0
-        for index in range(len(history) - 1, -1, -1):
-            if history[index].role in {"user", "assistant"}:
-                n_preserved += 1
-                if n_preserved == self.MAX_PRESERVED_MESSAGES:
-                    preserve_start_index = index
-                    break
+    def type_check(simple: SimpleCompaction):
+        _: Compaction = simple
 
-        if n_preserved < self.MAX_PRESERVED_MESSAGES:
-            return history
 
-        to_compact = history[:preserve_start_index]
-        to_preserve = history[preserve_start_index:]
+class SimpleCompaction:
+    def __init__(self, max_preserved_messages: int = 2) -> None:
+        self.max_preserved_messages = max_preserved_messages
 
-        if not to_compact:
-            # Let's hope this won't exceed the context size limit
+    async def compact(self, messages: Sequence[Message], llm: LLM) -> Sequence[Message]:
+        compact_message, to_preserve = self.prepare(messages)
+        if compact_message is None:
             return to_preserve
 
-        # Convert history to string for the compact prompt
-        history_text = "\n\n".join(
-            f"## Message {i + 1}\nRole: {msg.role}\nContent: {msg.content}"
-            for i, msg in enumerate(to_compact)
-        )
-
-        # Build the compact prompt using string template
-        compact_template = Template(prompts.COMPACT)
-        compact_prompt = compact_template.substitute(CONTEXT=history_text)
-
-        # Create input message for compaction
-        compact_message = Message(role="user", content=compact_prompt)
-
         # Call kosong.step to get the compacted context
         # TODO: set max completion tokens
         logger.debug("Compacting context...")
@@ -100,8 +74,42 @@ async def compact(self, messages: Sequence[Message], llm: LLM) -> Sequence[Messa
         compacted_messages.extend(to_preserve)
         return compacted_messages
 
+    class PrepareResult(NamedTuple):
+        compact_message: Message | None
+        to_preserve: Sequence[Message]
 
-if TYPE_CHECKING:
+    def prepare(self, messages: Sequence[Message]) -> PrepareResult:
+        if not messages or self.max_preserved_messages <= 0:
+            return self.PrepareResult(compact_message=None, to_preserve=messages)
 
-    def type_check(simple: SimpleCompaction):
-        _: Compaction = simple
+        history = list(messages)
+        preserve_start_index = len(history)
+        n_preserved = 0
+        for index in range(len(history) - 1, -1, -1):
+            if history[index].role in {"user", "assistant"}:
+                n_preserved += 1
+                if n_preserved == self.max_preserved_messages:
+                    preserve_start_index = index
+                    break
+
+        if n_preserved < self.max_preserved_messages:
+            return self.PrepareResult(compact_message=None, to_preserve=messages)
+
+        to_compact = history[:preserve_start_index]
+        to_preserve = history[preserve_start_index:]
+
+        if not to_compact:
+            # Let's hope this won't exceed the context size limit
+            return self.PrepareResult(compact_message=None, to_preserve=to_preserve)
+
+        # Create input message for compaction
+        compact_message = Message(role="user", content=[])
+        for i, msg in enumerate(to_compact):
+            compact_message.content.append(
+                TextPart(text=f"## Message {i + 1}\nRole: {msg.role}\nContent:\n")
+            )
+            compact_message.content.extend(
+                part for part in msg.content if not isinstance(part, ThinkPart)
+            )
+        compact_message.content.append(TextPart(text="\n" + prompts.COMPACT))
+        return self.PrepareResult(compact_message=compact_message, to_preserve=to_preserve)

tests/test_simple_compaction.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from inline_snapshot import snapshot
+from kosong.message import Message, TextPart, ThinkPart
+
+import kimi_cli.prompts as prompts
+from kimi_cli.soul.compaction import SimpleCompaction
+
+
+def test_prepare_returns_original_when_not_enough_messages():
+    messages = [Message(role="user", content=[TextPart(text="Only one message")])]
+
+    result = SimpleCompaction(max_preserved_messages=2).prepare(messages)
+
+    assert result == snapshot(
+        SimpleCompaction.PrepareResult(
+            compact_message=None,
+            to_preserve=[Message(role="user", content=[TextPart(text="Only one message")])],
+        )
+    )
+
+
+def test_prepare_skips_compaction_with_only_preserved_messages():
+    messages = [
+        Message(role="user", content=[TextPart(text="Latest question")]),
+        Message(role="assistant", content=[TextPart(text="Latest reply")]),
+    ]
+
+    result = SimpleCompaction(max_preserved_messages=2).prepare(messages)
+
+    assert result == snapshot(
+        SimpleCompaction.PrepareResult(
+            compact_message=None,
+            to_preserve=[
+                Message(role="user", content=[TextPart(text="Latest question")]),
+                Message(role="assistant", content=[TextPart(text="Latest reply")]),
+            ],
+        )
+    )
+
+
+def test_prepare_builds_compact_message_and_preserves_tail():
+    messages = [
+        Message(role="system", content=[TextPart(text="System note")]),
+        Message(
+            role="user",
+            content=[TextPart(text="Old question"), ThinkPart(think="Hidden thoughts")],
+        ),
+        Message(role="assistant", content=[TextPart(text="Old answer")]),
+        Message(role="user", content=[TextPart(text="Latest question")]),
+        Message(role="assistant", content=[TextPart(text="Latest answer")]),
+    ]
+
+    result = SimpleCompaction(max_preserved_messages=2).prepare(messages)
+
+    assert result.compact_message == snapshot(
+        Message(
+            role="user",
+            content=[
+                TextPart(text="## Message 1\nRole: system\nContent:\n"),
+                TextPart(text="System note"),
+                TextPart(text="## Message 2\nRole: user\nContent:\n"),
+                TextPart(text="Old question"),
+                TextPart(text="## Message 3\nRole: assistant\nContent:\n"),
+                TextPart(text="Old answer"),
+                TextPart(text="\n" + prompts.COMPACT),
+            ],
+        )
+    )
+    assert result.to_preserve == snapshot(
+        [
+            Message(role="user", content=[TextPart(text="Latest question")]),
+            Message(role="assistant", content=[TextPart(text="Latest answer")]),
+        ]
+    )