Detach provider state when branching chats
Rebuild branched chats from their trimmed local history by dropping inherited Responses continuation, response-ID ownership, and cached Context Window data. Preserve structured response output metadata so local replay remains faithful while source and branch provider lifecycles stay independent.
Alessandro committed
Aug 26, 2026 at 01:31 UTC
bb5d305091dec0c8edbde6e7f8f3a7335a9719da
4 files changed
+149
-1
plugins/_chat_branching/AGENTS.md
+4
@@ -16,6 +16,9 @@
16
17
- Preserve UUID-based linking between log entries and history messages.
18
- Branched chats must include history only up to the selected message.
19
+- Detach inherited Responses continuation and response-ID ownership, while
20
+ retaining structured output metadata needed for local replay.
21
+- Clear cached Context Window state after trimming each cloned agent.
22
- Do not mutate the source chat while creating a branch.
23
24
## Work Guidance
@@ -24,6 +27,7 @@
27
28
## Verification
29
30
+- Run `conda run -n a0 pytest plugins/_chat_branching/tests`.
31
- Smoke-test branching from several message positions and confirm source chat remains unchanged.
32
33
## Child DOX Index
plugins/_chat_branching/README.md
+2
@@ -15,6 +15,8 @@ Adds a **Branch** button to every chat message. Clicking it clones the current c
15
- Serializes the source context → deserializes into a new context with a fresh ID.
16
- Walks log entries: keeps everything up to the selected `log_no`, discards the rest.
17
- Collects the IDs of kept entries and uses them to trim `history.messages` so log and history stay consistent.
18
+ - Detaches inherited provider-response state and clears the cached Context
19
+ Window so the branch rebuilds both from its trimmed history.
20
21
3. **Persist & refresh**
22
- Saves the branched chat immediately.
plugins/_chat_branching/api/branch_chat.py
+20
-1
@@ -1,5 +1,6 @@
1
import json
2
3
+from agent import Agent, AgentContext
4
from helpers.api import ApiHandler, Input, Output, Request, Response
5
from helpers.localization import Localization
6
from helpers.persist_chat import (
@@ -7,7 +8,20 @@ from helpers.persist_chat import (
8
_deserialize_context,
9
save_tmp_chat,
10
)
10
-from agent import AgentContext
11
+
12
+
13
+def _detach_response_ids(value: object) -> None:
14
+ if isinstance(value, dict):
15
+ metadata = value.get("metadata")
16
+ responses = metadata.get("responses") if isinstance(metadata, dict) else None
17
+ if isinstance(responses, dict):
18
+ responses.pop("response_id", None)
19
+ responses.pop("previous_response_id", None)
20
+ for nested in value.values():
21
+ _detach_response_ids(nested)
22
+ elif isinstance(value, list):
23
+ for nested in value:
24
+ _detach_response_ids(nested)
25
26
27
def _trim_history_json(history_json: str, kept_ids: set[str], after_cut_ids: set[str]) -> str:
@@ -71,6 +85,7 @@ def _trim_history_json(history_json: str, kept_ids: set[str], after_cut_ids: set
85
len(t.get("messages", [])) for t in hist["topics"] if not t.get("summary")
86
) + len(hist.get("current", {}).get("messages", []))
87
hist["counter"] = total
88
+ _detach_response_ids(hist)
89
90
return json.dumps(hist, ensure_ascii=False)
91
@@ -124,6 +139,10 @@ class BranchChat(ApiHandler):
139
ag["history"] = _trim_history_json(
140
ag.get("history", ""), kept_ids, after_cut_ids
141
)
142
+ agent_data = ag.get("data")
143
+ if isinstance(agent_data, dict):
144
+ agent_data.pop(Agent.DATA_NAME_RESPONSES_STATE, None)
145
+ agent_data.pop(Agent.DATA_NAME_CTX_WINDOW, None)
146
147
# Give the branch a distinguishable name
148
src_name = data.get("name") or "Chat"
plugins/_chat_branching/tests/test_branch_chat.py
new
+123
@@ -0,0 +1,123 @@
1
+import copy
2
+import json
3
+from pathlib import Path
4
+import sys
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+
9
+
10
+ROOT = Path(__file__).resolve().parents[3]
11
+if str(ROOT) not in sys.path:
12
+ sys.path.insert(0, str(ROOT))
13
+
14
+from helpers import state_monitor_integration
15
+from helpers.persist_chat import _collect_response_ids
16
+from plugins._chat_branching.api import branch_chat
17
+
18
+
19
+@pytest.mark.asyncio
20
+async def test_branch_rebuilds_provider_and_context_state_from_trimmed_history(
21
+ monkeypatch,
22
+):
23
+ history = json.dumps(
24
+ {
25
+ "_cls": "History",
26
+ "counter": 2,
27
+ "bulks": [],
28
+ "topics": [],
29
+ "current": {
30
+ "summary": "",
31
+ "messages": [
32
+ {
33
+ "id": "kept-message",
34
+ "content": "before",
35
+ "metadata": {
36
+ "responses": {
37
+ "response_id": "resp_kept",
38
+ "previous_response_id": "resp_previous",
39
+ "output_items": [{"type": "message"}],
40
+ }
41
+ },
42
+ },
43
+ {"id": "removed-message", "content": "after"},
44
+ ],
45
+ },
46
+ }
47
+ )
48
+ serialized = {
49
+ "id": "source-chat",
50
+ "name": "Source chat",
51
+ "log": {
52
+ "logs": [
53
+ {"no": 4, "id": "kept-message"},
54
+ {"no": 5, "id": "removed-message"},
55
+ ]
56
+ },
57
+ "agents": [
58
+ {
59
+ "history": history,
60
+ "data": {
61
+ "responses_state": {
62
+ "response_id": "resp_current",
63
+ "response_ids": ["resp_kept", "resp_current"],
64
+ },
65
+ "ctx_window": {"text": "source context"},
66
+ },
67
+ }
68
+ for _ in range(2)
69
+ ],
70
+ }
71
+
72
+ branched = []
73
+
74
+ monkeypatch.setattr(
75
+ branch_chat.AgentContext,
76
+ "get",
77
+ lambda context_id: object() if context_id == "source-chat" else None,
78
+ )
79
+ monkeypatch.setattr(
80
+ branch_chat,
81
+ "_serialize_context",
82
+ lambda _context: copy.deepcopy(serialized),
83
+ )
84
+
85
+ def deserialize(data):
86
+ branched.append(copy.deepcopy(data))
87
+ return SimpleNamespace(id="branch-chat")
88
+
89
+ monkeypatch.setattr(branch_chat, "_deserialize_context", deserialize)
90
+ monkeypatch.setattr(branch_chat, "save_tmp_chat", lambda _context: None)
91
+ monkeypatch.setattr(
92
+ state_monitor_integration,
93
+ "mark_dirty_all",
94
+ lambda **_kwargs: None,
95
+ )
96
+
97
+ result = await branch_chat.BranchChat.process(
98
+ None,
99
+ {"context": "source-chat", "log_no": 4},
100
+ None,
101
+ )
102
+
103
+ assert result["ctxid"] == "branch-chat"
104
+ assert len(branched) == 1
105
+ assert _collect_response_ids(branched[0]) == []
106
+ for agent_data in branched[0]["agents"]:
107
+ assert "ctx_window" not in agent_data["data"]
108
+ assert "responses_state" not in agent_data["data"]
109
+ trimmed = json.loads(agent_data["history"])
110
+ messages = trimmed["current"]["messages"]
111
+ assert [message["id"] for message in messages] == ["kept-message"]
112
+ responses = messages[0]["metadata"]["responses"]
113
+ assert "response_id" not in responses
114
+ assert "previous_response_id" not in responses
115
+ assert responses["output_items"] == [{"type": "message"}]
116
+
117
+ assert serialized["agents"][0]["data"]["responses_state"]["response_id"] == (
118
+ "resp_current"
119
+ )
120
+ original_message = json.loads(serialized["agents"][0]["history"])["current"][
121
+ "messages"
122
+ ][0]
123
+ assert original_message["metadata"]["responses"]["response_id"] == "resp_kept"