feat: Add chat branching plugin
keyboardstaff committed
Feb 24, 2026 at 21:42 UTC
4633ebdd82d57b3ca5e4cbada61a6becb9d54c6f
3 files changed
+137
plugins/chat_branching/api/branch_chat.py
new
+75
@@ -0,0 +1,75 @@
1
+from datetime import datetime
2
+
3
+from python.helpers.api import ApiHandler, Input, Output, Request, Response
4
+from python.helpers.persist_chat import (
5
+ _serialize_context,
6
+ _deserialize_context,
7
+ save_tmp_chat,
8
+)
9
+from agent import AgentContext
10
+
11
+
12
+class BranchChat(ApiHandler):
13
+ """Create a new chat branched from an existing chat at a specific log message."""
14
+
15
+ async def process(self, input: Input, request: Request) -> Output:
16
+ ctxid = input.get("context", "")
17
+ log_no = input.get("log_no") # LogItem.no from frontend
18
+
19
+ if not ctxid:
20
+ return Response("Missing context id", 400)
21
+ if log_no is None:
22
+ return Response("Missing log_no", 400)
23
+
24
+ context = AgentContext.get(ctxid)
25
+ if not context:
26
+ return Response("Context not found", 404)
27
+
28
+ # Serialize the source context
29
+ data = _serialize_context(context)
30
+
31
+ # Remove id so _deserialize_context generates a new one
32
+ del data["id"]
33
+
34
+ # 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.
38
+ src_logs = data["log"]["logs"]
39
+ cut_idx = None
40
+ for i, item in enumerate(src_logs):
41
+ if item["no"] == log_no:
42
+ cut_idx = i
43
+ break
44
+
45
+ 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.
49
+ if 0 <= log_no < len(src_logs):
50
+ cut_idx = log_no
51
+ else:
52
+ return Response("log_no not found in chat log", 400)
53
+
54
+ data["log"]["logs"] = src_logs[: cut_idx + 1]
55
+
56
+ # Give the branch a distinguishable name
57
+ src_name = data.get("name") or "Chat"
58
+ data["name"] = f"{src_name} (branch)"
59
+ data["created_at"] = datetime.now().isoformat()
60
+
61
+ # Deserialize into a brand-new context (new id, fresh agent config)
62
+ new_context = _deserialize_context(data)
63
+
64
+ # Persist immediately
65
+ save_tmp_chat(new_context)
66
+
67
+ # Notify all tabs
68
+ from python.helpers.state_monitor_integration import mark_dirty_all
69
+ mark_dirty_all(reason="plugins.chat_branching.BranchChat")
70
+
71
+ return {
72
+ "ok": True,
73
+ "ctxid": new_context.id,
74
+ "message": "Chat branched successfully.",
75
+ }
\ No newline at end of file
plugins/chat_branching/extensions/webui/set_messages_after_loop/inject-branch-buttons.js
new
+58
@@ -0,0 +1,58 @@
1
+// Chat Branching Plugin — injects a "branch" button into every message's action bar.
2
+// Runs as a set_messages_after_loop JS extension.
3
+
4
+import { createActionButton } from "/components/messages/action-buttons/simple-action-buttons.js";
5
+import { callJsonApi } from "/js/api.js";
6
+
7
+const BRANCH_ATTR = "data-branch-injected";
8
+const LOG_NO_ATTR = "data-log-no";
9
+
10
+/**
11
+ * default export called by callJsExtensions("set_messages_after_loop", context)
12
+ * context.messages is the raw log items array with { no, id, type, ... }
13
+ */
14
+export default async function injectBranchButtons(context) {
15
+ if (!context?.messages?.length) return;
16
+
17
+ // 1. Stamp every rendered element with its log "no" so the button can read it.
18
+ for (const msg of context.messages) {
19
+ const domId = msg.id || msg.no;
20
+ // message containers use id="message-{id}", process steps use id="process-step-{id}"
21
+ const el =
22
+ document.getElementById(`message-${domId}`) ||
23
+ document.getElementById(`process-step-${domId}`);
24
+ if (el) el.setAttribute(LOG_NO_ATTR, String(msg.no));
25
+ }
26
+
27
+ // 2. Find every action-button bar that hasn't been patched yet and append a branch btn.
28
+ const bars = document.querySelectorAll(`.step-action-buttons:not([${BRANCH_ATTR}])`);
29
+ for (const bar of bars) {
30
+ bar.setAttribute(BRANCH_ATTR, "1");
31
+
32
+ // Resolve the log no from the nearest stamped ancestor
33
+ const stamped = bar.closest(`[${LOG_NO_ATTR}]`);
34
+ if (!stamped) continue;
35
+ const logNo = Number(stamped.getAttribute(LOG_NO_ATTR));
36
+ if (Number.isNaN(logNo)) continue;
37
+
38
+ const btn = createActionButton("fork_right", "Branch chat", async () => {
39
+ const ctxid = globalThis.getContext?.();
40
+ if (!ctxid) throw new Error("No active chat");
41
+
42
+ const res = await callJsonApi("/plugins/chat_branching/branch_chat", {
43
+ context: ctxid,
44
+ log_no: logNo,
45
+ });
46
+
47
+ if (!res?.ok) throw new Error(res?.message || "Branch failed");
48
+
49
+ // Select the newly created branch chat
50
+ const chatsStore = Alpine.store("chats");
51
+ if (chatsStore) {
52
+ chatsStore.selectChat(res.ctxid);
53
+ }
54
+ });
55
+
56
+ if (btn) bar.appendChild(btn);
57
+ }
58
+}
\ No newline at end of file
plugins/chat_branching/plugin.yaml
new
+4
@@ -0,0 +1,4 @@
1
+name: Chat Branching
2
+description: Branch a chat from any message, creating a new chat with history up to that point.
3
+version: 1.0.0
4
+always_enabled: true
\ No newline at end of file