perf(kosong): buffer stream merges and avoid deep-copying every delta

[parthgupta9999]

LLM streams often emit tiny text/tool-argument chunks. Concatenating with str += on each merge is quadratic, and model_copy(deep=True) on every callback was needlessly expensive on long responses.

Related Issue

N/A — performance improvement on the streaming merge path in kosong.

Description

Problem

While streaming a model response, kosong merges consecutive deltas (text, thinking, and tool-call argument fragments) into a single part. That path was doing two expensive things on every chunk:

1. Quadratic string growth — TextPart, ThinkPart, ToolCall, and ToolCallPart used str += inside merge_in_place. Providers usually send very small chunks, so a long answer (or large tool args) kept copying the whole accumulated string over and over.

2. Deep-copy on every callback — _generate called part.model_copy(deep=True) before passing each delta to on_message_part. Most stream parts are just short immutable strings, so a full deep copy per token was unnecessary work on the event loop.

Both run on every streamed turn, so the cost grows with response length.

Fix

In packages/kosong/src/kosong/message.py:

• Stream merges now buffer chunks in a private _merge_buf list during merge_in_place.
• Added finalize_merge() to join the buffer once into the public field.
• _message_append calls finalize_merge() before the part is added to the message, so the final content stays the same as before.

In packages/kosong/src/kosong/_generate.py:

• Replaced model_copy(deep=True) with _copy_stream_part().
• Most parts use a shallow model_copy().
• For ToolCall, the nested FunctionBody is also copied so later argument merges do not mutate the copy already sent to the callback.

Impact

• Final merged messages and tool-call arguments are unchanged.
• Less CPU and fewer allocations while streaming long text or chunked tool arguments.
• UI still gets per-delta callbacks; only merge buffering and copy strategy changed.

---

[devin-ai-integration]

Devin Review found 1 potential issue.

[devin-ai-integration]

🔴 Recorded and non-interactive outputs keep only the first fragment of streamed text and tool arguments

Merged text and tool-call arguments are only buffered, never combined (merge_in_place at packages/kosong/src/kosong/message.py:96-98), until a separate flush step runs, but the wire's merge path never runs that step, so only the first streamed fragment is kept.
Impact: Persisted sessions, background-agent output, subagent output, and print/stream-json output show only the first small chunk of each streamed assistant message and each tool call's arguments instead of the full content.

Contract change to merge_in_place not propagated to wire/print callers

Previously merge_in_place immediately concatenated into the public field (self.text += ..., self.function.arguments += ...). Now it appends into a private _merge_buf (packages/kosong/src/kosong/message.py:96-98, :126-128, :233-242, :263-272) and the public field is only updated when finalize_merge() runs (:102-105, etc.).

The PR only added the finalize_merge() call inside _message_append (packages/kosong/src/kosong/_generate.py:133), which covers the kosong generation path. However, other callers of merge_in_place were not updated:

• WireSoulSide in src/kimi_cli/wire/__init__.py:88-106 builds a merged buffer via self._merge_buffer.merge_in_place(msg) and later flush() publishes self._merge_buffer via _send_merged without ever calling finalize_merge(). The published/persisted part therefore carries only the first delta's text/arguments (the rest sit unused in _merge_buf, which is a PrivateAttr excluded from model_dump). This merged queue feeds the wire.jsonl recorder (src/kimi_cli/wire/__init__.py:31), the print UI (src/kimi_cli/ui/print/visualize.py:183), the background agent runner (src/kimi_cli/background/agent_runner.py:165), and the subagent runner (src/kimi_cli/subagents/runner.py:404).

• src/kimi_cli/ui/print/visualize.py similarly calls merge_in_place (lines 30 and 75) and then serializes via model_dump_json/extract_text without calling finalize_merge().

The interactive shell UI is unaffected because it accumulates text itself (append, append_args_part) rather than using merge_in_place.

Fix: call finalize_merge() on the buffered part before it is read/serialized — e.g. in WireSoulSide.flush() before _send_merged, and in the print visualizer before building the Message.

Prompt for agents

The change to MergeableMixin.merge_in_place in packages/kosong/src/kosong/message.py made merged content deferred: merge_in_place now appends chunks into a private _merge_buf, and the public field (text/think/arguments/arguments_part) is only updated when finalize_merge() is called. The PR added finalize_merge() only in kosong's _generate._message_append, but there are other callers of merge_in_place in kimi-cli that read/serialize the public field without calling finalize_merge(), so they now lose everything except the first streamed chunk.

Affected callers:
1. src/kimi_cli/wire/__init__.py WireSoulSide.send/flush (lines ~88-106): it deep-copies the first message into self._merge_buffer, then merges subsequent messages via merge_in_place, and flush() publishes self._merge_buffer to the merged queue via _send_merged without finalizing. The merged queue feeds the wire.jsonl recorder, the print UI, the background agent runner, and the subagent runner. Fix: call self._merge_buffer.finalize_merge() inside flush() (after asserting it is a wire message and before _send_merged), so the merged/persisted part contains the fully concatenated text/arguments.
2. src/kimi_cli/ui/print/visualize.py: _merge_content (line ~30) and JsonPrinter's ToolCall merge (line ~75) call merge_in_place, then build a Message and call model_dump_json/extract_text without finalizing. Fix: call finalize_merge() on each buffered part before constructing/serializing the Message (e.g. in _flush_assistant_message, FinalOnlyTextPrinter.flush, and FinalOnlyJsonPrinter.flush).

Verify no other callers of merge_in_place exist that read the public field without finalize_merge().

---

Was this helpful? React with 👍 or 👎 to provide feedback.