fix(chat_branching): not trimming agent history
keyboardstaff committed
Mar 23, 2026 at 22:37 UTC
aafe7f4aa886ed32e10caa747fe83f8a26163818
1 file changed
+59
-7
plugins/_chat_branching/api/branch_chat.py
+59
-7
@@ -1,3 +1,4 @@
1
+import json
2
from datetime import datetime
3
4
from helpers.api import ApiHandler, Input, Output, Request, Response
@@ -8,6 +9,57 @@ from helpers.persist_chat import (
9
)
10
from agent import AgentContext
11
12
+# Log types that produce a history message (via hist_add_* in Agent)
13
+_HIST_LOG_TYPES = {"user", "agent", "tool", "warning"}
14
+
15
+
16
+def _count_history_msgs(logs, agent_no):
17
+ """Count log items that produced history messages for a specific agent."""
18
+ return sum(
19
+ 1 for item in logs
20
+ if item["type"] in _HIST_LOG_TYPES and item.get("agentno", 0) == agent_no
21
+ )
22
+
23
+
24
+def _trim_history_json(history_json, keep_count):
25
+ """Trim a serialized history JSON string to keep only the first
26
+ *keep_count* un-summarized messages. Summarized bulks/topics are
27
+ always preserved (they represent older history before any cut point)."""
28
+ if not history_json:
29
+ return history_json
30
+ hist = json.loads(history_json)
31
+
32
+ # Flatten messages from topics + current, count and trim
33
+ remaining = keep_count
34
+ trimmed_topics = []
35
+
36
+ for topic in hist.get("topics", []):
37
+ if topic.get("summary"):
38
+ # Summarized topic — keep whole, don't count
39
+ trimmed_topics.append(topic)
40
+ continue
41
+ msgs = topic.get("messages", [])
42
+ if remaining >= len(msgs):
43
+ trimmed_topics.append(topic)
44
+ remaining -= len(msgs)
45
+ elif remaining > 0:
46
+ topic["messages"] = msgs[:remaining]
47
+ trimmed_topics.append(topic)
48
+ remaining = 0
49
+ # else: drop entirely
50
+
51
+ hist["topics"] = trimmed_topics
52
+
53
+ # Trim current topic
54
+ current = hist.get("current", {})
55
+ if not current.get("summary"):
56
+ msgs = current.get("messages", [])
57
+ if remaining < len(msgs):
58
+ current["messages"] = msgs[:remaining]
59
+
60
+ hist["counter"] = keep_count
61
+ return json.dumps(hist, ensure_ascii=False)
62
+
63
64
class BranchChat(ApiHandler):
65
"""Create a new chat branched from an existing chat at a specific log message."""
@@ -32,9 +84,6 @@ class BranchChat(ApiHandler):
84
del data["id"]
85
86
# Trim log entries: keep only items up to and including log_no.
35
- # _serialize_log uses log.logs[-LOG_SIZE:], so the serialized "no"
36
- # values may start above 0 for long chats. We match against the
37
- # original "no" field that each LogItem.output() emits.
87
src_logs = data["log"]["logs"]
88
cut_idx = None
89
for i, item in enumerate(src_logs):
@@ -43,15 +92,18 @@ class BranchChat(ApiHandler):
92
break
93
94
if cut_idx is None:
46
- # Fallback: log_no might already be a 0-based index within the
47
- # serialized array (e.g. after a reload where _deserialize_log
48
- # resets "no" to sequential). Accept if within bounds.
95
if 0 <= log_no < len(src_logs):
96
cut_idx = log_no
97
else:
98
return Response("log_no not found in chat log", 400)
99
54
- data["log"]["logs"] = src_logs[: cut_idx + 1]
100
+ kept_logs = src_logs[: cut_idx + 1]
101
+ data["log"]["logs"] = kept_logs
102
+
103
+ # Trim agent history to match the log cut point
104
+ for ag in data.get("agents", []):
105
+ msg_count = _count_history_msgs(kept_logs, ag["number"])
106
+ ag["history"] = _trim_history_json(ag.get("history", ""), msg_count)
107
108
# Give the branch a distinguishable name
109
src_name = data.get("name") or "Chat"