@samitouri / QOSami_Agents_Zero / commits / 1bd5bc01

refactor(chat_branching): ID-based log ↔ history linking for precise branch trimming

keyboardstaff committed Mar 24, 2026 at 03:04 UTC 1bd5bc01d66cab860864769b830f3cac9f94ecfb

22 files changed +165 -99

agent.py

+20 -15

@@ -320,6 +320,7 @@ class UserMessage:

320 message: str

321 attachments: list[str] = field(default_factory=list[str])

322 system_message: list[str] = field(default_factory=list[str])

323 + id: str = ""

324

325

326 class LoopData:

@@ -466,18 +467,20 @@ class Agent:

467 self.loop_data.last_response == agent_response

468 ): # if assistant_response is the same as last message in history, let him know

469 # Append the assistant's response to the history

469 - self.hist_add_ai_response(agent_response)

470 + log_item = self.loop_data.params_temporary.get("log_item_generating")

471 + self.hist_add_ai_response(agent_response, id=log_item.id if log_item else "")

472 # Append warning message to the history

473 warning_msg = self.read_prompt("fw.msg_repeat.md")

472 - self.hist_add_warning(message=warning_msg)

474 + wmsg = self.hist_add_warning(message=warning_msg)

475 PrintStyle(font_color="orange", padding=True).print(

476 warning_msg

477 )

476 - self.context.log.log(type="warning", content=warning_msg)

478 + self.context.log.log(type="warning", content=warning_msg, id=wmsg.id)

479

480 else: # otherwise proceed with tool

481 # Append the assistant's response to the history

480 - self.hist_add_ai_response(agent_response)

482 + log_item = self.loop_data.params_temporary.get("log_item_generating")

483 + self.hist_add_ai_response(agent_response, id=log_item.id if log_item else "")

484 # process tools requested in agent message

485 tools_result = await self.process_tools(agent_response)

486 if tools_result: # final response of message loop available

@@ -637,7 +640,7 @@ class Agent:

640

641 @extension.extensible

642 def hist_add_message(

640 - self, ai: bool, content: history.MessageContent, tokens: int = 0

643 + self, ai: bool, content: history.MessageContent, tokens: int = 0, id: str = ""

644 ):

645 self.last_message = datetime.now(timezone.utc)

646 # Allow extensions to process content before adding to history

@@ -646,7 +649,7 @@ class Agent:

649 "hist_add_before", self, content_data=content_data, ai=ai

650 )

651 return self.history.add_message(

649 - ai=ai, content=content_data["content"], tokens=tokens

652 + ai=ai, content=content_data["content"], tokens=tokens, id=id

653 )

654

655 @extension.extensible

@@ -674,30 +677,31 @@ class Agent:

677 content = {k: v for k, v in content.items() if v}

678

679 # add to history

677 - msg = self.hist_add_message(False, content=content) # type: ignore

680 + msg = self.hist_add_message(False, content=content, id=message.id) # type: ignore

681 self.last_user_message = msg

682 return msg

683

684 @extension.extensible

682 - def hist_add_ai_response(self, message: str):

685 + def hist_add_ai_response(self, message: str, id: str = ""):

686 self.loop_data.last_response = message

687 content = self.parse_prompt("fw.ai_response.md", message=message)

685 - return self.hist_add_message(True, content=content)

688 + return self.hist_add_message(True, content=content, id=id)

689

690 @extension.extensible

688 - def hist_add_warning(self, message: history.MessageContent):

691 + def hist_add_warning(self, message: history.MessageContent, id: str = ""):

692 content = self.parse_prompt("fw.warning.md", message=message)

690 - return self.hist_add_message(False, content=content)

693 + return self.hist_add_message(False, content=content, id=id)

694

695 @extension.extensible

696 def hist_add_tool_result(self, tool_name: str, tool_result: str, **kwargs):

697 + msg_id = kwargs.pop("id", "")

698 data = {

699 "tool_name": tool_name,

700 "tool_result": tool_result,

701 **kwargs,

702 }

703 extension.call_extensions_sync("hist_add_tool_result", self, data=data)

700 - return self.hist_add_message(False, content=data)

704 + return self.hist_add_message(False, content=data, id=msg_id)

705

706 def concat_messages(

707 self, messages

@@ -927,18 +931,19 @@ class Agent:

931 error_detail = (

932 f"Tool '{raw_tool_name}' not found or could not be initialized."

933 )

930 - self.hist_add_warning(error_detail)

934 + wmsg = self.hist_add_warning(error_detail)

935 PrintStyle(font_color="red", padding=True).print(error_detail)

936 self.context.log.log(

933 - type="warning", content=f"{self.agent_name}: {error_detail}"

937 + type="warning", content=f"{self.agent_name}: {error_detail}", id=wmsg.id

938 )

939 else:

940 warning_msg_misformat = self.read_prompt("fw.msg_misformat.md")

937 - self.hist_add_warning(warning_msg_misformat)

941 + wmsg = self.hist_add_warning(warning_msg_misformat)

942 PrintStyle(font_color="red", padding=True).print(warning_msg_misformat)

943 self.context.log.log(

944 type="warning",

945 content=f"{self.agent_name}: Message misformat, no valid tool request found.",

946 + id=wmsg.id,

947 )

948

949 @extension.extensible

api/api_message.py

+4 -1

@@ -1,5 +1,6 @@

1 import base64

2 import os

3 +import uuid

4 from datetime import datetime, timedelta

5 from agent import AgentContext, UserMessage, AgentContextType

6 from helpers.api import ApiHandler, Request, Response

@@ -134,15 +135,17 @@ class ApiMessage(ApiHandler):

135 PrintStyle(font_color="white", padding=False).print(f"- {filename}")

136

137 # Add user message to chat history so it's visible in the UI

138 + msg_id = str(uuid.uuid4())

139 context.log.log(

140 type="user",

141 heading="",

142 content=message,

143 kvps={"attachments": attachment_filenames},

144 + id=msg_id,

145 )

146

147 # Send message to agent

145 - task = context.communicate(UserMessage(message=message, attachments=attachment_paths))

148 + task = context.communicate(UserMessage(message=message, attachments=attachment_paths, id=msg_id))

149 result = await task.result()

150

151 # Clean up expired chats

api/message.py

+1 -1

@@ -68,4 +68,4 @@ class Message(ApiHandler):

68 # Log to console and UI using helper function

69 mq.log_user_message(context, message, attachment_paths, message_id)

71 - return context.communicate(UserMessage(message=message, attachments=attachment_paths)), context

71 + return context.communicate(UserMessage(message=message, attachments=attachment_paths, id=message_id or "")), context

extensions/python/_functions/agent/Agent/handle_exception/end/_50_handle_repairable_exception.py

+2 -2

@@ -17,9 +17,9 @@ class HandleRepairableException(Extension):

17 if isinstance(data["exception"], RepairableException):

18 msg = {"message": errors.format_error(data["exception"])}

19 await extension.call_extensions_async("error_format", agent=self.agent, msg=msg)

20 - self.agent.hist_add_warning(msg["message"])

20 + wmsg = self.agent.hist_add_warning(msg["message"])

21 PrintStyle(font_color="red", padding=True).print(msg["message"])

22 - self.agent.context.log.log(type="warning", content=msg["message"])

22 + self.agent.context.log.log(type="warning", content=msg["message"], id=wmsg.id)

23 data["exception"] = None

extensions/python/agent_init/_10_initial_message.py

+2 -1

@@ -28,7 +28,7 @@ class InitialMessage(Extension):

28 self.agent.loop_data = LoopData(user_message=None)

30 # Add the message to history as an AI response

31 - self.agent.hist_add_ai_response(initial_message)

31 + msg = self.agent.hist_add_ai_response(initial_message)

33 # json parse the message, get the tool_args text

34 initial_message_json = json.loads(initial_message)

@@ -40,4 +40,5 @@ class InitialMessage(Extension):

40 content=initial_message_text,

41 finished=True,

42 update_progress="none",

43 + id=msg.id,

44 )

extensions/python/before_main_llm_call/_10_log_for_stream.py

@@ -5,6 +5,7 @@ import asyncio

5 from helpers.log import LogItem

6 from helpers import log

7 import math

8 +import uuid

11 class LogForStream(Extension):

@@ -19,6 +20,7 @@ class LogForStream(Extension):

20 self.agent.context.log.log(

21 type="agent",

22 heading=build_default_heading(self.agent),

23 + id=str(uuid.uuid4()),

24 )

25 )

extensions/python/response_stream/_20_live_response.py

@@ -30,10 +30,14 @@ class LiveResponse(Extension):

31 # create log message and store it in loop data temporary params

32 if "log_item_response" not in loop_data.params_temporary:

33 + # Share id with the agent log item so branching covers the response bubble

34 + gen_item = loop_data.params_temporary.get("log_item_generating")

35 + shared_id = gen_item.id if gen_item and gen_item.id else ""

36 loop_data.params_temporary["log_item_response"] = (

37 self.agent.context.log.log(

38 type="response",

39 heading=f"icon://chat {self.agent.agent_name}: Responding",

40 + id=shared_id,

41 )

42 )

helpers/history.py

+9 -6

@@ -4,6 +4,7 @@ from collections import OrderedDict

4 from collections.abc import Mapping

5 import json

6 import math

7 +import uuid

8 from typing import Coroutine, Literal, TypedDict, cast, Union, Dict, List, Any

9 from helpers import messages, tokens, settings, call_llm

10 from enum import Enum

@@ -81,7 +82,8 @@ class Record:

84 class Message(Record):

84 - def __init__(self, ai: bool, content: MessageContent, tokens: int = 0):

85 + def __init__(self, ai: bool, content: MessageContent, tokens: int = 0, id: str = ""):

86 + self.id = id or str(uuid.uuid4())

87 self.ai = ai

88 self.content = content

89 self.summary: str = ""

@@ -115,6 +117,7 @@ class Message(Record):

117 def to_dict(self):

118 return {

119 "_cls": "Message",

120 + "id": self.id,

121 "ai": self.ai,

122 "content": self.content,

123 "summary": self.summary,

@@ -124,7 +127,7 @@ class Message(Record):

127 @staticmethod

128 def from_dict(data: dict, history: "History"):

129 content = data.get("content", "Content lost")

127 - msg = Message(ai=data["ai"], content=content)

130 + msg = Message(ai=data["ai"], content=content, id=data.get("id", ""))

131 msg.summary = data.get("summary", "")

132 msg.tokens = data.get("tokens", 0)

133 return msg

@@ -143,9 +146,9 @@ class Topic(Record):

146 return sum(msg.get_tokens() for msg in self.messages)

147

148 def add_message(

146 - self, ai: bool, content: MessageContent, tokens: int = 0

149 + self, ai: bool, content: MessageContent, tokens: int = 0, id: str = ""

150 ) -> Message:

148 - msg = Message(ai=ai, content=content, tokens=tokens)

151 + msg = Message(ai=ai, content=content, tokens=tokens, id=id)

152 self.messages.append(msg)

153 return msg

154

@@ -332,10 +335,10 @@ class History(Record):

335 return self.current.get_tokens()

336

337 def add_message(

335 - self, ai: bool, content: MessageContent, tokens: int = 0

338 + self, ai: bool, content: MessageContent, tokens: int = 0, id: str = ""

339 ) -> Message:

340 self.counter += 1

338 - return self.current.add_message(ai, content=content, tokens=tokens)

341 + return self.current.add_message(ai, content=content, tokens=tokens, id=id)

342

343 def new_topic(self):

344 if self.current.messages:

helpers/mcp_handler.py

+3 -1

@@ -102,11 +102,13 @@ class MCPTool(Tool):

102 """MCP Tool wrapper"""

103

104 def get_log_object(self) -> LogItem:

105 + import uuid

106 return self.agent.context.log.log(

107 type="mcp",

108 heading=f"icon://extension {self.agent.agent_name}: Using MCP tool '{self.name}'",

109 content="",

110 kvps={"tool_name": self.name, **self.args},

111 + id=str(uuid.uuid4()),

112 )

113

114 async def execute(self, **kwargs: Any):

@@ -198,7 +200,7 @@ class MCPTool(Tool):

200

201 final_text_for_agent = raw_tool_response

202

201 - self.agent.hist_add_tool_result(self.name, final_text_for_agent)

203 + self.agent.hist_add_tool_result(self.name, final_text_for_agent, id=self.log.id if self.log else "")

204 (

205 PrintStyle(

206 font_color="#1B4F72", background_color="white", padding=True, bold=True

helpers/message_queue.py

+6 -4

@@ -153,8 +153,9 @@ def send_message(context: "AgentContext", item: dict, source: str = " (from queu

153

154 message = item.get("text", "")

155 attachments = item.get("attachments", [])

156 - log_user_message(context, message, attachments, source=source)

157 - context.communicate(UserMessage(message, attachments))

156 + msg_id = str(uuid.uuid4())

157 + log_user_message(context, message, attachments, message_id=msg_id, source=source)

158 + context.communicate(UserMessage(message, attachments, id=msg_id))

159

160

161 def send_next(context: "AgentContext") -> bool:

@@ -183,6 +184,7 @@ def send_all_aggregated(context: "AgentContext") -> int:

184 text = "\n\n---\n\n".join(i["text"] for i in items if i["text"])

185 attachments = [a for i in items for a in i.get("attachments", [])]

186

186 - log_user_message(context, text, attachments, source=" (queued batch)")

187 - context.communicate(UserMessage(text, attachments))

187 + msg_id = str(uuid.uuid4())

188 + log_user_message(context, text, attachments, message_id=msg_id, source=" (queued batch)")

189 + context.communicate(UserMessage(text, attachments, id=msg_id))

190 return len(items)

helpers/task_scheduler.py

+4 -2

@@ -882,19 +882,21 @@ class TaskScheduler:

882 task_prompt = f"## Task:\n{current_task.prompt}"

883

884 # Log the message with message_id and attachments

885 + msg_id = str(uuid.uuid4())

886 context.log.log(

887 type="user",

888 heading="",

889 content=task_prompt,

890 kvps={"attachments": attachment_filenames},

890 - id=str(uuid.uuid4()),

891 + id=msg_id,

892 )

893

894 agent.hist_add_user_message(

895 UserMessage(

896 message=task_prompt,

897 system_message=[current_task.system_prompt],

897 - attachments=attachment_filenames))

898 + attachments=attachment_filenames,

899 + id=msg_id))

900

901 # Persist after setting up the context but before running the agent

902 # This ensures the task context is saved and can be found by polling

helpers/tool.py

+4 -2

@@ -53,17 +53,19 @@ class Tool:

54 async def after_execution(self, response: Response, **kwargs):

55 text = sanitize_string(response.message.strip())

56 - self.agent.hist_add_tool_result(self.name, text, **(response.additional or {}))

56 + self.agent.hist_add_tool_result(self.name, text, id=self.log.id, **(response.additional or {}))

57 PrintStyle(font_color="#1B4F72", background_color="white", padding=True, bold=True).print(f"{self.agent.agent_name}: Response from tool '{self.name}'")

58 PrintStyle(font_color="#85C1E9").print(text)

59 self.log.update(content=text)

61 def get_log_object(self):

62 + import uuid

63 + pre_id = str(uuid.uuid4())

64 if self.method:

65 heading = f"icon://construction {self.agent.agent_name}: Using tool '{self.name}:{self.method}'"

66 else:

67 heading = f"icon://construction {self.agent.agent_name}: Using tool '{self.name}'"

66 - return self.agent.context.log.log(type="tool", heading=heading, content="", kvps=self.args, _tool_name=self.name)

68 + return self.agent.context.log.log(type="tool", heading=heading, content="", kvps=self.args, _tool_name=self.name, id=pre_id)

70 def nice_key(self, key:str):

71 words = key.split('_')

plugins/_chat_branching/README.md

+17 -15

@@ -4,26 +4,28 @@ Create a new chat from any existing point in a conversation.

5 ## What It Does

7 -This plugin adds an API endpoint that clones an existing chat context, trims its log history up to a selected message, gives the new chat a `(branch)` suffix, persists it immediately, and refreshes connected tabs so the new branch appears in the UI.

7 +Adds a **Branch** button to every chat message. Clicking it clones the current chat up to that message, creating a new conversation you can continue independently.

9 -## Main Behavior

9 +## How It Works

11 -- **Clone context**

12 - - Serializes the current chat context and deserializes it into a brand-new context with a new ID.

13 -- **Trim history**

14 - - Keeps log entries only up to and including the selected `log_no`.

15 - - Includes a fallback for cases where reloaded logs use sequential array indexes.

16 -- **Persist immediately**

17 - - Saves the newly created branched chat to temporary chat storage.

18 -- **Refresh UI state**

19 - - Marks state dirty for all tabs after the branch is created.

11 +1. **ID-based log ↔ history linking**

12 + Every `LogItem` and `history.Message` share a UUID generated at creation time. The branch button is only shown on messages that carry this ID.

13 +

14 +2. **Clone & trim**

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 +

19 +3. **Persist & refresh**

20 + - Saves the branched chat immediately.

21 + - Marks UI state dirty so all connected tabs see the new branch.

23 ## Entry Points

23 -- **API**

24 - - `api/branch_chat.py` implements the branching operation.

25 -- **Extensions**

26 - - `extensions/` contains integration glue for making the feature available in the app.

25 +| Path | Purpose |

26 +|---|---|

27 +| `api/branch_chat.py` | API endpoint — clone, trim, persist |

28 +| `extensions/webui/set_messages_after_loop/inject-branch-buttons.js` | Injects the Branch button into each message DOM element |

30 ## Plugin Metadata

plugins/_chat_branching/api/branch_chat.py

+56 -36

@@ -9,55 +9,69 @@ from helpers.persist_chat import (

9 )

10 from agent import AgentContext

12 -# Log types that produce a history message (via hist_add_* in Agent)

13 -_HIST_LOG_TYPES = {"user", "agent", "tool", "warning"}

13 +def _trim_history_json(history_json: str, kept_ids: set[str], after_cut_ids: set[str]) -> str:

14 + """Trim a serialized history JSON to keep only messages that appear

15 + before the branch cut point.

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 - )

17 + *kept_ids*: IDs of log items up to and including the cut point.

18 + *after_cut_ids*: IDs of log items after the cut point.

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)."""

20 + Walk messages in order. A message is kept while the running state

21 + is "keep". The state flips to "drop" the first time we encounter

22 + a message whose id is in *after_cut_ids*. Messages whose id does

23 + not appear in any log set (unpaired) inherit the current state.

24 + Summarized topics/bulks are always preserved.

25 + """

26 if not history_json:

27 return history_json

30 - hist = json.loads(history_json)

32 - # Flatten messages from topics + current, count and trim

33 - remaining = keep_count

29 + hist = json.loads(history_json)

30 + keep = True # running state

31 +

32 + def filter_messages(messages: list[dict]) -> list[dict]:

33 + nonlocal keep

34 + result = []

35 + for msg in messages:

36 + mid = msg.get("id", "")

37 + if mid and mid in after_cut_ids:

38 + keep = False

39 + elif mid and mid in kept_ids:

40 + keep = True

41 + # else: unpaired – inherit current state

42 + if keep:

43 + result.append(msg)

44 + return result

45 +

46 + # Bulks are already summarized old history – always keep

47 + # Topics: keep summarized ones; filter unsummarized

48 trimmed_topics = []

35 -

49 for topic in hist.get("topics", []):

50 if topic.get("summary"):

38 - # Summarized topic — keep whole, don't count

51 trimmed_topics.append(topic)

52 continue

41 - msgs = topic.get("messages", [])

42 - if remaining >= len(msgs):

53 + msgs = filter_messages(topic.get("messages", []))

54 + if msgs:

55 + topic["messages"] = msgs

56 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 -

57 + if not keep:

58 + break

59 hist["topics"] = trimmed_topics

53 - # Trim current topic

61 + # Current topic

62 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]

63 + if not current.get("summary") and keep:

64 + current["messages"] = filter_messages(current.get("messages", []))

65 + elif not keep:

66 + current["messages"] = []

67 + hist["current"] = current

68 +

69 + # Recount

70 + total = sum(

71 + len(t.get("messages", [])) for t in hist["topics"] if not t.get("summary")

72 + ) + len(hist.get("current", {}).get("messages", []))

73 + hist["counter"] = total

60 - hist["counter"] = keep_count

75 return json.dumps(hist, ensure_ascii=False)

@@ -98,12 +112,18 @@ class BranchChat(ApiHandler):

112 return Response("log_no not found in chat log", 400)

113

114 kept_logs = src_logs[: cut_idx + 1]

115 + after_logs = src_logs[cut_idx + 1 :]

116 data["log"]["logs"] = kept_logs

117

103 - # Trim agent history to match the log cut point

118 + # Build ID sets for history trimming

119 + kept_ids = {item["id"] for item in kept_logs if item.get("id")}

120 + after_cut_ids = {item["id"] for item in after_logs if item.get("id")}

121 +

122 + # Trim each agent's history using ID matching

123 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)

124 + ag["history"] = _trim_history_json(

125 + ag.get("history", ""), kept_ids, after_cut_ids

126 + )

127

128 # Give the branch a distinguishable name

129 src_name = data.get("name") or "Chat"

plugins/_chat_branching/extensions/webui/set_messages_after_loop/inject-branch-buttons.js

+1 -1

@@ -9,7 +9,7 @@ export default async function injectBranchButtons(context) {

9 if (!context?.results?.length) return;

11 for (const { args, result } of context.results) {

12 - if (!result?.element || args.no == null) continue;

12 + if (!result?.element || args.no == null || !args.id) continue;

14 const logNo = args.no;

15 for (const bar of result.element.querySelectorAll(".step-action-buttons")) {

plugins/_code_execution/tools/code_execution_tool.py

+3 -1

@@ -71,11 +71,13 @@ class CodeExecution(Tool):

71 return Response(message=response, break_loop=False)

73 def get_log_object(self):

74 + import uuid

75 return self.agent.context.log.log(

76 type="code_exe",

77 heading=self.get_heading(),

78 content="",

79 kvps=self.args,

80 + id=str(uuid.uuid4()),

81 )

83 def get_heading(self, text: str = ""):

@@ -86,7 +88,7 @@ class CodeExecution(Tool):

88 return f"icon://terminal {session_text}{truncate_text_string(text, 200)}"

90 async def after_execution(self, response, **kwargs):

89 - self.agent.hist_add_tool_result(self.name, response.message, **(response.additional or {}))

91 + self.agent.hist_add_tool_result(self.name, response.message, id=self.log.id if self.log else "", **(response.additional or {}))

93 async def prepare_state(self, cfg: dict, reset=False, session: int | None = None):

94 self.state: State | None = self.agent.get_data("_cet_state")

plugins/_code_execution/tools/input.py

+3 -2

@@ -19,7 +19,8 @@ class Input(Tool):

19 return await cet.execute(**args)

21 def get_log_object(self):

22 - return self.agent.context.log.log(type="code_exe", heading=f"icon://keyboard {self.agent.agent_name}: Using tool '{self.name}'", content="", kvps=self.args)

22 + import uuid

23 + return self.agent.context.log.log(type="code_exe", heading=f"icon://keyboard {self.agent.agent_name}: Using tool '{self.name}'", content="", kvps=self.args, id=str(uuid.uuid4()))

25 async def after_execution(self, response, **kwargs):

25 - self.agent.hist_add_tool_result(self.name, response.message, **(response.additional or {}))

\ No newline at end of file

26 + self.agent.hist_add_tool_result(self.name, response.message, id=self.log.id if self.log else "", **(response.additional or {}))

\ No newline at end of file

plugins/_email_integration/helpers/handler.py

+7 -2

@@ -8,6 +8,7 @@ import asyncio

8 import base64

9 import json

10 import os

11 +import uuid

13 from agent import Agent, AgentContext, AgentContextType, UserMessage

14 from helpers import guids, plugins, files, runtime

@@ -285,11 +286,13 @@ async def _start_new_chat(agent: Agent, handler_cfg: dict, msg: InboundMessage):

286 user_msg = _build_user_message(agent, msg, handler_cfg)

287 system_ctx = agent.read_prompt("fw.email.system_context.md")

288

288 - mq.log_user_message(context, user_msg, msg.attachments or [], source=" (email)")

289 + msg_id = str(uuid.uuid4())

290 + mq.log_user_message(context, user_msg, msg.attachments or [], message_id=msg_id, source=" (email)")

291 context.communicate(UserMessage(

292 message=user_msg,

293 system_message=[system_ctx],

294 attachments=msg.attachments,

295 + id=msg_id,

296 ))

297

298 PrintStyle.success(f"Email: new chat {context.id} for '{msg.subject}' from {msg.sender}")

@@ -322,10 +325,12 @@ async def _route_to_chat(

325 context.data[disp.CTX_EMAIL_REFERENCES] = " ".join(refs_list)

326

327 user_msg = _build_user_message(agent, msg, handler_cfg)

325 - mq.log_user_message(context, user_msg, msg.attachments or [], source=" (email)")

328 + msg_id = str(uuid.uuid4())

329 + mq.log_user_message(context, user_msg, msg.attachments or [], message_id=msg_id, source=" (email)")

330 context.communicate(UserMessage(

331 message=user_msg,

332 attachments=msg.attachments,

333 + id=msg_id,

334 ))

335

336 save_tmp_chat(context)

plugins/_error_retry/extensions/python/_functions/agent/Agent/handle_exception/end/_80_retry_critical_exception.py

+4 -1

@@ -35,10 +35,13 @@ class RetryCriticalException(Extension):

35 self.agent.set_data(DATA_NAME_COUNTER, counter + 1)

37 error_message = errors.format_error(exception)

38 + import uuid as _uuid

39 + msg_id = str(_uuid.uuid4())

40 self.agent.context.log.log(

41 type="warning",

42 heading="Critical error occurred, retrying...",

43 content=error_message,

44 + id=msg_id,

45 )

46 PrintStyle(font_color="orange", padding=True).print(

47 "Critical error occurred, retrying..."

@@ -48,7 +51,7 @@ class RetryCriticalException(Extension):

51 agent_facing_error = self.agent.read_prompt(

52 "fw.msg_critical_error.md", error_message=error_message

53 )

51 - self.agent.hist_add_warning(message=agent_facing_error)

54 + self.agent.hist_add_warning(message=agent_facing_error, id=msg_id)

55 PrintStyle(font_color="orange", padding=True).print(agent_facing_error)

57 data["exception"] = None

plugins/_infection_check/helpers/checker.py

+5 -1

@@ -328,11 +328,14 @@ class InfectionChecker:

328 return "terminate", "Max clarifications exceeded.", "\n\n".join(cot_parts)

329

330 def _do_terminate(self, agent: "Agent", detail: str, cot: str):

331 + import uuid as _uuid

332 content = cot or detail or "Malicious behavior detected."

333 + msg_id = str(_uuid.uuid4())

334 agent.context.log.log(

335 type="warning",

336 heading="Infection check: TERMINATED",

337 content=content,

338 + id=msg_id,

339 )

340

341 # Replace last AI message with a blocked marker

@@ -341,7 +344,8 @@ class InfectionChecker:

344 if msgs and msgs[-1].ai:

345 msgs.pop()

346 agent.history.add_message(

344 - ai=True, content="[BLOCKED] Response terminated by security policy."

347 + ai=True, content="[BLOCKED] Response terminated by security policy.",

348 + id=msg_id,

349 )

350 except Exception:

351 pass

plugins/_telegram_integration/helpers/handler.py

+6 -3

@@ -246,10 +246,12 @@ async def handle_message(message: TgMessage, bot_name: str, bot_cfg: dict):

246 body=text,

247 )

248

249 - mq.log_user_message(context, user_msg, attachments, source=" (telegram)")

249 + msg_id = str(uuid.uuid4())

250 + mq.log_user_message(context, user_msg, attachments, message_id=msg_id, source=" (telegram)")

251 context.communicate(UserMessage(

252 message=user_msg,

253 attachments=attachments,

254 + id=msg_id,

255 ))

256

257 save_tmp_chat(context)

@@ -298,8 +300,9 @@ async def handle_callback_query(query: CallbackQuery, bot_name: str, bot_cfg: di

300 body=f"[Button pressed: {text}]",

301 )

302

301 - mq.log_user_message(context, user_msg, [], source=" (telegram)")

302 - context.communicate(UserMessage(message=user_msg))

303 + msg_id = str(uuid.uuid4())

304 + mq.log_user_message(context, user_msg, [], message_id=msg_id, source=" (telegram)")

305 + context.communicate(UserMessage(message=user_msg, id=msg_id))

306 save_tmp_chat(context)

307

308

tools/vision_load.py

+2 -2

@@ -75,7 +75,7 @@ class VisionLoad(Tool):

75 f"Skipped images (max {self._get_max_embeds()} loaded at a time according to model configuration):\n{skipped_summary}"

76 )

77 if self.images_dict:

78 - self.agent.hist_add_tool_result(self.name, summary)

78 + self.agent.hist_add_tool_result(self.name, summary, id=self.log.id if self.log else "")

79 for path, image in self.images_dict.items():

80 if image:

81 content.append(

@@ -97,7 +97,7 @@ class VisionLoad(Tool):

97 False, content=msg, tokens=TOKENS_ESTIMATE * len(content)

98 )

99 else:

100 - self.agent.hist_add_tool_result(self.name, summary if self.skipped_paths else "No images processed")

100 + self.agent.hist_add_tool_result(self.name, summary if self.skipped_paths else "No images processed", id=self.log.id if self.log else "")

101

102 # print and log short version

103 message = (