diff --git a/plugins/chat_branching/extensions/webui/set_messages_after_loop/inject-branch-buttons.js b/plugins/chat_branching/extensions/webui/set_messages_after_loop/inject-branch-buttons.js
index eeb377bf4..80568b613 100644
--- a/plugins/chat_branching/extensions/webui/set_messages_after_loop/inject-branch-buttons.js
+++ b/plugins/chat_branching/extensions/webui/set_messages_after_loop/inject-branch-buttons.js
@@ -1,41 +1,22 @@
 // Chat Branching Plugin — injects a "branch" button into every message's action bar.
 // Runs as a set_messages_after_loop JS extension.
+//
+// Uses the unified handler output: context.results is an array of
+// { args: { no, id, type, … }, result: { element, … } } from each setMessage() call.
 
 import { createActionButton } from "/components/messages/action-buttons/simple-action-buttons.js";
 import { callJsonApi } from "/js/api.js";
 
-const LOG_NO_ATTR = "data-log-no";
-
-/**
- * default export called by callJsExtensions("set_messages_after_loop", context)
- * context.messages is the raw log items array with { no, id, type, ... }
- */
 export default async function injectBranchButtons(context) {
-  if (!context?.messages?.length) return;
+  if (!context?.results?.length) return;
 
-  // 1. Stamp every rendered element with its log "no" so the button can read it.
-  for (const msg of context.messages) {
-    const domId = msg.id || msg.no;
-    // message containers use id="message-{id}", process steps use id="process-step-{id}"
-    const el =
-      document.getElementById(`message-${domId}`) ||
-      document.getElementById(`process-step-${domId}`);
-    if (el) el.setAttribute(LOG_NO_ATTR, String(msg.no));
-  }
+  for (const { args, result } of context.results) {
+    if (!result?.element || args.no == null) continue;
 
-  // 2. Inject branch button into every action bar.
-  //    Core renderers clear and rebuild action-button children on each update,
-  //    so we check for the actual button element instead of a marker attribute.
-  const bars = document.querySelectorAll(".step-action-buttons");
-  for (const bar of bars) {
-    if (bar.querySelector(".action-fork_right")) continue;
-
-    // Resolve the log no from the nearest stamped ancestor
-    const stamped = bar.closest(`[${LOG_NO_ATTR}]`);
-    if (!stamped) continue;
-    const logNo = Number(stamped.getAttribute(LOG_NO_ATTR));
-    if (Number.isNaN(logNo)) continue;
+    const bar = result.element.querySelector(".step-action-buttons");
+    if (!bar || bar.querySelector(".action-fork_right")) continue;
 
+    const logNo = args.no;
     const btn = createActionButton("fork_right", "Branch chat", async () => {
       const ctxid = globalThis.getContext?.();
       if (!ctxid) throw new Error("No active chat");
diff --git a/plugins/chat_branching/plugin.yaml b/plugins/chat_branching/plugin.yaml
index bb382ed22..7ca67b278 100644
--- a/plugins/chat_branching/plugin.yaml
+++ b/plugins/chat_branching/plugin.yaml
@@ -1,4 +1,3 @@
 name: Chat Branching
 description: Branch a chat from any message, creating a new chat with history up to that point.
-version: 1.0.0
-always_enabled: true
\ No newline at end of file
+version: 1.0.0
\ No newline at end of file