From 951db80a07f6f1c5bc694557cd448c883b3237b3 Mon Sep 17 00:00:00 2001 From: Michele Papalini Date: Mon, 7 Sep 2026 15:43:55 +0200 Subject: [PATCH] fix: keep tool_calls/tool_result pairs atomic when trimming conversation history Signed-off-by: Michele Papalini --- .../standard/plugins/context/conversation.py | 48 +++++------------ .../standard/plugins/context/token_budget.py | 25 +-------- .../standard/plugins/context/tool_pairing.py | 51 +++++++++++++++++++ 3 files changed, 67 insertions(+), 57 deletions(-) create mode 100644 library-standard/src/mas/library/standard/plugins/context/tool_pairing.py diff --git a/library-standard/src/mas/library/standard/plugins/context/conversation.py b/library-standard/src/mas/library/standard/plugins/context/conversation.py index a71596d1..7182a843 100644 --- a/library-standard/src/mas/library/standard/plugins/context/conversation.py +++ b/library-standard/src/mas/library/standard/plugins/context/conversation.py @@ -7,6 +7,7 @@ import logging from typing import Any, Callable +from mas.library.standard.plugins.context.tool_pairing import group_exchanges, skip_tool_group from mas.runtime.contracts.context_manager_contract import ContextManagerContract _log = logging.getLogger(__name__) @@ -29,13 +30,20 @@ def manage_history( return past if len(past) <= self.max_messages: return past - evicted = len(past) - self.max_messages + + tail = list(past) + while len(tail) > self.max_messages: + n = skip_tool_group(tail, 0) + if n >= len(tail): + break + del tail[:n] + _log.debug( "StackConversation: evicted %d message(s), keeping last %d", - evicted, - self.max_messages, + len(past) - len(tail), + len(tail), ) - return past[-self.max_messages :] + return tail class SlidingWindowConversation(ContextManagerContract): @@ -54,21 +62,7 @@ def manage_history( if not past: return past - exchanges: list[list[dict[str, Any]]] = [] - i = 0 - while i < len(past): - msg = past[i] - if msg.get("role") == "user": - exchange: list[dict[str, Any]] = [msg] - if i + 1 < len(past) and past[i + 1].get("role") == "assistant": - exchange.append(past[i + 1]) - i += 2 - else: - i += 1 - exchanges.append(exchange) - else: - exchanges.append([msg]) - i += 1 + exchanges = group_exchanges(past) if len(exchanges) <= self.max_turns: return past @@ -123,21 +117,7 @@ def manage_history( if self._estimate_tokens(past) <= effective: return past - exchanges: list[list[dict[str, Any]]] = [] - i = 0 - while i < len(past): - msg = past[i] - if msg.get("role") == "user": - exchange: list[dict[str, Any]] = [msg] - if i + 1 < len(past) and past[i + 1].get("role") == "assistant": - exchange.append(past[i + 1]) - i += 2 - else: - i += 1 - exchanges.append(exchange) - else: - exchanges.append([msg]) - i += 1 + exchanges = group_exchanges(past) if len(exchanges) <= self.keep_turns: return past diff --git a/library-standard/src/mas/library/standard/plugins/context/token_budget.py b/library-standard/src/mas/library/standard/plugins/context/token_budget.py index ce30ff71..7544eb93 100644 --- a/library-standard/src/mas/library/standard/plugins/context/token_budget.py +++ b/library-standard/src/mas/library/standard/plugins/context/token_budget.py @@ -6,6 +6,8 @@ from typing import Any +from mas.library.standard.plugins.context.tool_pairing import skip_tool_group as _skip_tool_group + def estimate_tokens(messages: list[dict[str, Any]]) -> int: total = 0 @@ -19,29 +21,6 @@ def estimate_tokens(messages: list[dict[str, Any]]) -> int: return total // 4 + len(messages) * 4 -def _skip_tool_group(tail: list[dict[str, Any]], start: int) -> int: - """Return how many messages to drop starting at *start* to keep tool pairs intact. - - If ``tail[start]`` is an assistant message with ``tool_calls``, we must also - drop every subsequent ``tool`` response that references one of those calls. - If ``tail[start]`` is an orphaned ``tool`` message, drop it too. - """ - msg = tail[start] - if msg.get("role") == "assistant" and msg.get("tool_calls"): - call_ids = {c.get("id") for c in msg["tool_calls"] if c.get("id")} - count = 1 - while start + count < len(tail): - nxt = tail[start + count] - if nxt.get("role") == "tool" and nxt.get("tool_call_id") in call_ids: - count += 1 - else: - break - return count - if msg.get("role") == "tool": - return 1 - return 1 - - def trim_messages_to_budget( messages: list[dict[str, Any]], *, diff --git a/library-standard/src/mas/library/standard/plugins/context/tool_pairing.py b/library-standard/src/mas/library/standard/plugins/context/tool_pairing.py new file mode 100644 index 00000000..44ed97b6 --- /dev/null +++ b/library-standard/src/mas/library/standard/plugins/context/tool_pairing.py @@ -0,0 +1,51 @@ +# Copyright (c) 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Shared helpers for trimming message lists without splitting tool_calls/tool pairs.""" + +from __future__ import annotations + +from typing import Any + + +def skip_tool_group(messages: list[dict[str, Any]], start: int) -> int: + """Return how many messages to drop starting at *start* to keep tool pairs intact. + + If ``messages[start]`` is an assistant message with ``tool_calls``, every + subsequent ``tool`` response that references one of those calls is dropped + with it. An orphaned ``tool`` message is dropped alone. + """ + msg = messages[start] + if msg.get("role") == "assistant" and msg.get("tool_calls"): + call_ids = {c.get("id") for c in msg["tool_calls"] if c.get("id")} + count = 1 + while start + count < len(messages): + nxt = messages[start + count] + if nxt.get("role") == "tool" and nxt.get("tool_call_id") in call_ids: + count += 1 + else: + break + return count + return 1 + + +def group_exchanges(past: list[dict[str, Any]]) -> list[list[dict[str, Any]]]: + """Group messages into user/assistant exchanges, keeping tool_calls/tool pairs atomic.""" + exchanges: list[list[dict[str, Any]]] = [] + i = 0 + while i < len(past): + msg = past[i] + if msg.get("role") == "user": + exchange = [msg] + i += 1 + if i < len(past) and past[i].get("role") == "assistant": + exchange.append(past[i]) + i += 1 + n = skip_tool_group(past, i - 1) - 1 # trailing tool results, if any + exchange.extend(past[i : i + n]) + i += n + exchanges.append(exchange) + else: + n = skip_tool_group(past, i) + exchanges.append(past[i : i + n]) + i += n + return exchanges