From ecd6974e3bb1a2085ae917ec0223f39ae99fc3ba Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Wed, 5 Aug 2026 01:15:06 +0800 Subject: [PATCH 01/14] Fix tool_call_id loss in unified memory message round-trip --- ms_agent/memory/unified/orchestrator.py | 37 +++++++++++++------------ tests/memory/test_unified_memory.py | 25 +++++++++++++++++ 2 files changed, 45 insertions(+), 17 deletions(-) diff --git a/ms_agent/memory/unified/orchestrator.py b/ms_agent/memory/unified/orchestrator.py index 000a643be..789096a88 100644 --- a/ms_agent/memory/unified/orchestrator.py +++ b/ms_agent/memory/unified/orchestrator.py @@ -12,6 +12,9 @@ from ms_agent.llm.utils import Message from ms_agent.memory.base import Memory +# Single canonical deserializer, shared with ContextAssembler. A second local +# copy previously drifted and silently dropped `tool_call_id` / `name`. +from ms_agent.session.context_assembler import _dicts_to_messages from ms_agent.utils.logger import get_logger from .config import MemoryConfig from .protocols import MemoryBackend, MemoryEntry @@ -181,6 +184,15 @@ def _default_base_dir(mc: MemoryConfig, config: Any) -> None: def _messages_to_dicts(messages: List[Message]) -> List[Dict[str, Any]]: + """Serialize for ``backend.inject()`` — must be lossless. + + This round-trip (``run()``: messages -> dicts -> inject -> messages) runs on + EVERY turn, so any field dropped here is dropped from the live LLM context, + not just from storage. ``tool_call_id`` in particular is mandatory on the + wire for ``role='tool'`` rows: losing it makes OpenAI-compatible providers + reject the request with ``missing field 'tool_call_id'``. Keep this the + exact inverse of ``_dicts_to_messages`` below. + """ result: List[Dict[str, Any]] = [] for m in messages: if isinstance(m, dict): @@ -189,24 +201,15 @@ def _messages_to_dicts(messages: List[Message]) -> List[Dict[str, Any]]: d: Dict[str, Any] = {'role': m.role, 'content': m.content or ''} if m.tool_calls: d['tool_calls'] = m.tool_calls + if m.tool_call_id: + d['tool_call_id'] = m.tool_call_id + if m.name: + d['name'] = m.name + if m.reasoning_content: + d['reasoning_content'] = m.reasoning_content + if m.reasoning_signature: + d['reasoning_signature'] = m.reasoning_signature result.append(d) else: result.append({'role': 'user', 'content': str(m)}) return result - - -def _dicts_to_messages(dicts: List[Dict[str, Any]]) -> List[Message]: - result: List[Message] = [] - for d in dicts: - if isinstance(d, Message): - result.append(d) - elif isinstance(d, dict): - result.append( - Message( - role=d.get('role', 'user'), - content=d.get('content', ''), - tool_calls=d.get('tool_calls'), - )) - else: - result.append(Message(role='user', content=str(d))) - return result diff --git a/tests/memory/test_unified_memory.py b/tests/memory/test_unified_memory.py index 2063b5949..e8391d797 100644 --- a/tests/memory/test_unified_memory.py +++ b/tests/memory/test_unified_memory.py @@ -262,6 +262,31 @@ def test_dicts_to_messages_preserves_tool_fields(self): assert msgs[1].tool_call_id == "call_1" assert msgs[1].name == "search" + def test_memory_orchestrator_round_trip_preserves_tool_fields(self): + """The unified-memory round-trip runs on EVERY turn, in-memory, before + the LLM call — so a field dropped there is dropped from the live + request. This previously regressed independently of the assembler + (the test above imported the assembler's copy and never covered it), + and OpenAI-compatible providers answered with + ``missing field 'tool_call_id'``. + """ + from ms_agent.memory.unified.orchestrator import (_dicts_to_messages, + _messages_to_dicts) + from ms_agent.llm.utils import Message + + msgs = [ + Message(role="assistant", content="", + tool_calls=[{"id": "call_1", "type": "function"}]), + Message(role="tool", content="result", + tool_call_id="call_1", name="search"), + ] + out = _dicts_to_messages(_messages_to_dicts(msgs)) + assert out[1].tool_call_id == "call_1" + assert out[1].name == "search" + # And it must actually reach the wire: to_dict_clean() drops falsy + # values, so a lost id disappears from the payload entirely. + assert out[1].to_dict_clean()["tool_call_id"] == "call_1" + # ═══════════════════════════════════════════════════════════════════════ # 2. ViewStrategies From f8ce706129b49b1c77e0ad84afff3c345f27acfd Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Wed, 5 Aug 2026 01:15:50 +0800 Subject: [PATCH 02/14] Pair tool results with pending calls when tool_call_id is missing in openai-compat transport --- ms_agent/llm/transport/openai_compat.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/ms_agent/llm/transport/openai_compat.py b/ms_agent/llm/transport/openai_compat.py index 3e7779c47..d25316419 100644 --- a/ms_agent/llm/transport/openai_compat.py +++ b/ms_agent/llm/transport/openai_compat.py @@ -171,6 +171,14 @@ def _format_input_message(self, cache_indice = max(cache_indices) if cache_indices else None openai_messages = [] + # Order-matched fallback for a tool row that reaches us without its id. + # OpenAI-compatible gateways reject such a message outright + # ("missing field `tool_call_id`"), killing the whole turn, so pair it + # with the preceding assistant turn's calls instead. The Anthropic + # transport has carried the same guard for a while; this keeps the two + # symmetric. Note `to_dict_clean()` omits falsy values, so a None/'' id + # disappears from the dict entirely rather than arriving as None. + pending_tool_ids: List[str] = [] for idx, message in enumerate(messages): if isinstance(message, Message): if isinstance(message.content, str): @@ -207,6 +215,22 @@ def _format_input_message(self, and not content): formatted_message['content'] = None + role = formatted_message.get('role') + if role == 'assistant': + pending_tool_ids = [ + tc.get('id') for tc in (formatted_message.get('tool_calls') + or []) if isinstance(tc, dict) + and tc.get('id') + ] + elif role == 'tool' and not formatted_message.get('tool_call_id'): + if pending_tool_ids: + formatted_message['tool_call_id'] = pending_tool_ids.pop(0) + else: + logger.warning( + 'tool message has no tool_call_id and no preceding ' + 'assistant tool_calls to match it against; the provider ' + 'will likely reject this request') + openai_messages.append(formatted_message) return openai_messages From a7b86ec557005db1be8f72d716a433bea22cb6ca Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Wed, 5 Aug 2026 01:15:52 +0800 Subject: [PATCH 03/14] Seal errored rounds so resume consumes the next prompt instead of replaying --- ms_agent/agent/llm_agent.py | 65 +++++++++++++++++++++++-------- tests/agent/test_partial_round.py | 64 ++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 17 deletions(-) diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index 2dd2b905a..a7364b731 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -59,8 +59,15 @@ _INTERRUPTED_TOOL_RESULT = '[Interrupted: tool execution was cancelled]' -def build_partial_round_records(rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Turn an interrupted round's in-memory rows into protocol-valid log records. +def build_partial_round_records( + rows: List[Dict[str, Any]], + marker: str = 'interrupted') -> List[Dict[str, Any]]: + """Turn an unfinished round's in-memory rows into protocol-valid log records. + + ``marker`` is the boolean key stamped on every produced record, naming *why* + the round ended early: ``interrupted`` (user Stop / cancellation) or + ``errored`` (the turn raised). UIs badge these differently, so an API + failure must not be sealed as an interruption. ``rows`` are the ``_msg_to_dict`` serializations of ``messages[pre_step_len:]`` at cancellation time. Rules (each keeps replay valid on BOTH transports): @@ -79,7 +86,7 @@ def build_partial_round_records(rows: List[Dict[str, Any]]) -> List[Dict[str, An with no content and no kept calls gets the neutral placeholder content so the turn reads as closed (an empty assistant block would be rejected on Anthropic replay, and a dangling user row would be re-answered on resume). - - every row is flagged ``interrupted: true`` — an extra key that survives in + - every row is flagged ``: true`` — an extra key that survives in the log for UI replay but is filtered out of the LLM context rebuild. """ present_results = { @@ -89,7 +96,7 @@ def build_partial_round_records(rows: List[Dict[str, Any]]) -> List[Dict[str, An records: List[Dict[str, Any]] = [] for row in rows: rec = dict(row) - rec['interrupted'] = True + rec[marker] = True if rec.get('role') != 'assistant': records.append(rec) continue @@ -130,14 +137,14 @@ def build_partial_round_records(rows: List[Dict[str, Any]]) -> List[Dict[str, An 'tool_call_id': tc.get('id'), 'name': tc.get('tool_name', ''), 'is_error': True, - 'interrupted': True, + marker: True, }) if not records: records.append({ 'role': 'assistant', 'content': INTERRUPTED_PLACEHOLDER, 'content_placeholder': True, - 'interrupted': True, + marker: True, }) return records @@ -1980,16 +1987,24 @@ def _msg_to_dict(msg: Message) -> Dict[str, Any]: d['tokens'] = prompt_tokens + completion_tokens return d - def _persist_partial_round(self, messages: List[Message], - pre_step_len: int) -> None: - """Seal an interrupted round into the SessionLog (best-effort). - - Called from run_loop's cancellation handler, where the normal - round-boundary persistence can no longer run. Serializes the round's - in-memory rows and appends the protocol-repaired records built by - :func:`build_partial_round_records`. Synchronous file I/O only (safe - inside a cancelled task); never raises — sealing must not break the - cancellation unwind. + def _persist_partial_round(self, + messages: List[Message], + pre_step_len: int, + marker: str = 'interrupted') -> None: + """Seal an unfinished round into the SessionLog (best-effort). + + Called from run_loop's cancellation handler AND its exception handler, + where the normal round-boundary persistence can no longer run. + Serializes the round's in-memory rows and appends the protocol-repaired + records built by :func:`build_partial_round_records`. Synchronous file + I/O only (safe inside a cancelled task); never raises — sealing must not + break the unwind. + + Sealing is not cosmetic: without it the log tail stays on a ``user`` or + ``tool`` row, and the next resume rebuilds that same round and re-calls + the model instead of reading the user's new prompt (see run_loop's + ``load_cache`` guard). ``marker`` records why the round ended — + ``interrupted`` for Stop, ``errored`` for a raised turn. """ if self.session_log is None: return @@ -2018,7 +2033,7 @@ def _persist_partial_round(self, messages: List[Message], rows = [ self._msg_to_dict(msg) for msg in messages[pre_step_len:] ] - for record in build_partial_round_records(rows): + for record in build_partial_round_records(rows, marker=marker): self.session_log.append(record) except Exception: logger.warning('persist partial round failed', exc_info=True) @@ -2034,6 +2049,11 @@ async def run_loop(self, messages: Union[List[Message], str], Args: messages: Input prompt string or list of Message objects. """ + # Bound BEFORE the try so the exception handler can always read it: + # setup (agent/LLM construction, credential resolution) raises well + # before the round loop, and an unbound local there would turn a clean + # provider error into an UnboundLocalError. None = no round to seal. + pre_step_len: Optional[int] = None try: self.max_chat_round = getattr(self.config, 'max_chat_round', LLMAgent.DEFAULT_MAX_CHAT_ROUND) @@ -2266,6 +2286,17 @@ async def run_loop(self, messages: Union[List[Message], str], import traceback logger.warning(traceback.format_exc()) + # Seal the round FIRST. A raised turn leaves the log tail on a + # `user`/`tool` row, and run_loop's resume guard only skips the + # model call when the tail is `assistant` — so the next request + # rebuilds this same round, re-sends the identical failing context, + # and never reaches after_tool_call (which is what drains the queued + # prompt). The user sees the turn "loop" while every resend is + # silently discarded. Sealing closes the round, so the rebuilt agent + # blocks on read_prompt and consumes the new message instead. + if pre_step_len is not None: + self._persist_partial_round( + messages, pre_step_len, marker='errored') if self._event_sink is not None: # A run_loop turn-abort is non-recoverable (the turn produced no # usable assistant output); mark it so live == persisted/replay. diff --git a/tests/agent/test_partial_round.py b/tests/agent/test_partial_round.py index 75d17d478..b0281dff8 100644 --- a/tests/agent/test_partial_round.py +++ b/tests/agent/test_partial_round.py @@ -185,3 +185,67 @@ def test_reasoning_duration_finalized_when_cancelled_mid_reasoning(): import time dur = _persisted_reasoning_duration(time.monotonic() - 3.0, None) assert dur is not None and dur >= 2 + + +# --- sealing a round that RAISED (not a user Stop) ----------------------- + + +def test_errored_marker_replaces_interrupted(): + # run_loop's exception handler seals with marker='errored'. UIs badge + # `interrupted` as a user Stop, so an API failure must not borrow that key. + records = build_partial_round_records([], marker='errored') + assert records == [{ + 'role': 'assistant', + 'content': INTERRUPTED_PLACEHOLDER, + 'content_placeholder': True, + 'errored': True, + }] + assert 'interrupted' not in records[0] + + +def test_errored_marker_applies_to_synthesized_tool_results(): + records = build_partial_round_records( + [{'role': 'assistant', 'content': '', + 'tool_calls': [{'id': 'c1', 'arguments': '{}', 'tool_name': 'grep'}]}], + marker='errored') + assert [r['role'] for r in records] == ['assistant', 'tool'] + assert all(r.get('errored') is True for r in records) + assert not any('interrupted' in r for r in records) + + +def test_error_seal_closes_dangling_tool_tail(): + """The regression that made failed turns look like an infinite loop. + + A raised turn used to leave the log tail on a `tool` row. run_loop's resume + guard only skips the model call when the tail is `assistant`, so the next + request replayed the same round with the same failing context and never + drained the queued prompt — every resend was silently dropped. + """ + from ms_agent.agent.llm_agent import LLMAgent + + class _Log: + + def __init__(self, rows): + self.rows = list(rows) + + def append(self, rec): + self.rows.append(rec) + return len(self.rows) + + # Tail as persisted by the round that then blew up on the next LLM call. + stub = LLMAgent.__new__(LLMAgent) + stub.session_log = _Log([ + {'role': 'user', 'content': '搜一下今天的热点新闻'}, + {'role': 'assistant', 'content': '', + 'tool_calls': [{'id': 'c1', 'arguments': '{}', 'tool_name': 'search'}]}, + {'role': 'tool', 'content': 'results', 'tool_call_id': 'c1'}, + ]) + stub._reasoning_started_at = None + stub._last_reasoning_duration = None + assert stub.session_log.rows[-1]['role'] == 'tool' # the poisoned state + + # step() raised before yielding anything -> empty segment. + LLMAgent._persist_partial_round(stub, [], 0, marker='errored') + + assert stub.session_log.rows[-1]['role'] == 'assistant' + assert stub.session_log.rows[-1]['errored'] is True From 0686305ffe91df4158748a1f6909680924a7e6cf Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Wed, 5 Aug 2026 01:16:18 +0800 Subject: [PATCH 04/14] Skip LLM call retries for non-retryable 4xx client errors --- ms_agent/agent/llm_agent.py | 9 +- ms_agent/utils/__init__.py | 2 +- ms_agent/utils/llm_utils.py | 61 ++++++++++++- tests/utils/test_retry_classification.py | 109 +++++++++++++++++++++++ 4 files changed, 175 insertions(+), 6 deletions(-) create mode 100644 tests/utils/test_retry_classification.py diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index a7364b731..03d30a0fb 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -40,7 +40,8 @@ ReasoningDelta, ReasoningEnded, ReasoningStarted, ToolCallCompleted, ToolCallStarted, TurnCompleted, UsageInfo) -from ms_agent.utils import async_retry, read_history, save_history +from ms_agent.utils import (async_retry, is_retryable_error, read_history, + save_history) from ms_agent.utils.constants import DEFAULT_TAG, DEFAULT_USER from ms_agent.utils.logger import get_logger from ms_agent.utils.snapshot import take_snapshot @@ -1535,7 +1536,11 @@ def _append_task_notifications(self, messages.append(Message(role='user', content=body)) return messages - @async_retry(max_attempts=Agent.retry_count, delay=1.0) + # retry_if: a hard 4xx (bad payload, content filter, auth) is a verdict on + # the request, not a transient fault — retrying it 5× only adds ~40s of + # backoff before the same failure surfaces. + @async_retry( + max_attempts=Agent.retry_count, delay=1.0, retry_if=is_retryable_error) async def step( self, messages: List[Message] ) -> AsyncGenerator[List[Message], Any]: # type: ignore diff --git a/ms_agent/utils/__init__.py b/ms_agent/utils/__init__.py index 32655bd93..afa87e941 100644 --- a/ms_agent/utils/__init__.py +++ b/ms_agent/utils/__init__.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .llm_utils import async_retry, retry +from .llm_utils import async_retry, is_retryable_error, retry from .logger import get_logger from .prompt import get_fact_retrieval_prompt from .utils import (assert_package_exist, enhance_error, read_history, diff --git a/ms_agent/utils/llm_utils.py b/ms_agent/utils/llm_utils.py index 4f800bd45..7d0611bc7 100644 --- a/ms_agent/utils/llm_utils.py +++ b/ms_agent/utils/llm_utils.py @@ -1,8 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import functools +import re import time -from typing import Any, AsyncGenerator, Callable, Tuple, Type, TypeVar, Union +from typing import (Any, AsyncGenerator, Callable, Optional, Tuple, Type, + TypeVar, Union) from .logger import get_logger @@ -10,6 +12,47 @@ T = TypeVar('T') +# 4xx statuses that CAN succeed on a resend: request timeout, lock conflict, +# too-early, and rate limiting. Every other 4xx is a verdict on the request +# itself — resending the identical payload just buys the same rejection. +_RETRYABLE_CLIENT_STATUSES = frozenset({408, 409, 425, 429}) + +# Providers surface the status inconsistently: the openai/anthropic SDKs expose +# `status_code`, while gateways often only put it in the message +# ("APIError: <400> ...", "Error code: 400 - {...}"). +_STATUS_IN_TEXT = re.compile(r'<(\d{3})>|(?:error\s+)?code[:=]?\s*(\d{3})\b', + re.IGNORECASE) + + +def http_status_of(exc: BaseException) -> Optional[int]: + """Best-effort HTTP status for a provider exception, or None if unknown.""" + for attr in ('status_code', 'http_status', 'status'): + value = getattr(exc, attr, None) + if isinstance(value, int) and 100 <= value <= 599: + return value + match = _STATUS_IN_TEXT.search(str(exc)) + if match: + status = int(match.group(1) or match.group(2)) + if 100 <= status <= 599: + return status + return None + + +def is_retryable_error(exc: BaseException) -> bool: + """Whether resending the same request could plausibly succeed. + + Unknown/None status keeps the historical behaviour (retry), so transient + network faults are unaffected. Only an identifiable, non-transient 4xx + fails fast — previously a hard 400 (bad payload, content filter, missing + field) burned all five attempts plus 15s of backoff before surfacing. + """ + status = http_status_of(exc) + if status is None: + return True + if 400 <= status < 500: + return status in _RETRYABLE_CLIENT_STATUSES + return True + def retry(max_attempts: int = 3, delay: float = 1.0, @@ -54,8 +97,14 @@ def async_retry(max_attempts: int = 3, delay: float = 1.0, backoff_factor: float = 2.0, exceptions: Union[Type[Exception], Tuple[Type[Exception], - ...]] = Exception): - """Retry doing something""" + ...]] = Exception, + retry_if: Optional[Callable[[BaseException], bool]] = None): + """Retry doing something. + + ``retry_if`` short-circuits the loop for exceptions that cannot succeed on + a resend (see :func:`is_retryable_error`); the exception is still raised, + just without burning the remaining attempts and their backoff. + """ def decorator(func: Callable[..., T]) -> Callable[..., T]: @@ -73,6 +122,12 @@ async def wrapper(*args, **kwargs) -> AsyncGenerator[T, Any]: import traceback logger.warning(traceback.format_exc()) last_exception = e + if retry_if is not None and not retry_if(e): + logger.error( + f'{func.__name__} failed unrecoverably on attempt ' + f'{attempt}/{max_attempts}; not retrying. ' + f'Exception message: {e}') + break if attempt < max_attempts: logger.warning( f'Attempt {attempt}/{max_attempts} fails: {func.__name__}. ' diff --git a/tests/utils/test_retry_classification.py b/tests/utils/test_retry_classification.py new file mode 100644 index 000000000..79cc3c969 --- /dev/null +++ b/tests/utils/test_retry_classification.py @@ -0,0 +1,109 @@ +"""async_retry's ``retry_if`` gate: don't burn attempts on unrecoverable 4xx. + +A hard 400 (bad payload, missing field, content filter) is a verdict on the +request — resending the identical body five times only adds ~15s of backoff +before the same failure surfaces. Transient faults must keep retrying. +""" +import asyncio + +import pytest +from ms_agent.utils import async_retry, is_retryable_error +from ms_agent.utils.llm_utils import http_status_of + + +class _ApiError(Exception): + """Stand-in for a provider SDK error, with or without ``status_code``.""" + + def __init__(self, message, status_code=None): + super().__init__(message) + if status_code is not None: + self.status_code = status_code + + +@pytest.mark.parametrize( + 'exc, expected', + [ + # Real strings observed in production, where the status is only in text. + (_ApiError('APIError: <400> InternalError.Algo.DataInspectionFailed: ' + 'Output data may contain inappropriate content.'), 400), + (_ApiError("Error code: 400 - {'error': {'message': 'missing field " + "`tool_call_id`'}}"), 400), + # Structured SDK errors. + (_ApiError('boom', 429), 429), + (_ApiError('boom', 503), 503), + # Nothing status-like. + (ConnectionError('Connection reset by peer'), None), + ], +) +def test_http_status_extraction(exc, expected): + assert http_status_of(exc) == expected + + +@pytest.mark.parametrize('status', [400, 401, 403, 404, 422]) +def test_client_errors_are_not_retryable(status): + assert is_retryable_error(_ApiError('boom', status)) is False + + +@pytest.mark.parametrize('status', [408, 409, 425, 429, 500, 502, 503]) +def test_transient_and_server_errors_stay_retryable(status): + assert is_retryable_error(_ApiError('boom', status)) is True + + +@pytest.mark.parametrize( + 'exc', + [ConnectionError('reset'), + TimeoutError('timed out'), + Exception('something odd')]) +def test_unknown_errors_keep_retrying(exc): + # Conservative default: an unidentifiable failure behaves as before. + assert is_retryable_error(exc) is True + + +def _run(fn): + async def drive(): + async for _ in fn(): + pass + + with pytest.raises(Exception): + asyncio.run(drive()) + + +def test_unrecoverable_error_uses_exactly_one_attempt(): + calls = [] + + @async_retry(max_attempts=5, delay=10.0, retry_if=is_retryable_error) + async def boom(): + calls.append(1) + raise _ApiError('APIError: <400> DataInspectionFailed') + yield # pragma: no cover - makes this an async generator + + _run(boom) + # delay=10.0 would make a retrying implementation take >=10s; one attempt + # also proves no backoff was slept. + assert len(calls) == 1 + + +def test_transient_error_still_exhausts_attempts(): + calls = [] + + @async_retry(max_attempts=3, delay=0.01, retry_if=is_retryable_error) + async def flaky(): + calls.append(1) + raise ConnectionError('reset') + yield # pragma: no cover + + _run(flaky) + assert len(calls) == 3 + + +def test_without_retry_if_behaviour_is_unchanged(): + calls = [] + + @async_retry(max_attempts=3, delay=0.01) + async def boom(): + calls.append(1) + raise _ApiError('APIError: <400> DataInspectionFailed') + yield # pragma: no cover + + _run(boom) + assert len(calls) == 3 From 977d16b9f4fddb96557c19669f631e39be2f31fe Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Wed, 5 Aug 2026 01:16:22 +0800 Subject: [PATCH 05/14] Dedupe identical per-round error records in SessionLog --- ms_agent/session/session_log.py | 13 ++++++++++ tests/memory/test_unified_memory.py | 37 +++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/ms_agent/session/session_log.py b/ms_agent/session/session_log.py index 84f922aab..8b8e4c6aa 100644 --- a/ms_agent/session/session_log.py +++ b/ms_agent/session/session_log.py @@ -103,7 +103,20 @@ def record_error(self, event: Dict[str, Any]) -> None: turn-level / API errors that must NOT go back to the model (tool-call errors stay as ordinary ``role="tool"`` messages instead). ``event`` typically carries ``message``, ``error_type``, ``recoverable``, ``round``. + + Idempotent per (round, message): a round that is retried or replayed + must not stack identical records. A wedged turn used to append one per + attempt, growing a log of dozens of copies of the same failure and + making replay unreadable. Dedup needs a ``round`` to key on; without one + every call is recorded (callers that omit it are opting out). """ + rnd = event.get('round') + if rnd is not None: + message = event.get('message') + for prior in self.get_errors(): + if (prior.get('round') == rnd + and prior.get('message') == message): + return seq = self._next_seq() record = { "_type": "error", diff --git a/tests/memory/test_unified_memory.py b/tests/memory/test_unified_memory.py index e8391d797..6f92cb347 100644 --- a/tests/memory/test_unified_memory.py +++ b/tests/memory/test_unified_memory.py @@ -1696,3 +1696,40 @@ def test_session_and_memory_pipeline(self): assert len(log.get_all_messages()) == 4 finally: loop.close() + + +class TestSessionLogErrorDedup: + """record_error is idempotent per (round, message). + + A wedged turn used to append one identical record per retry/replay attempt + (observed: five copies of the same 400, all round 15), making replay + unreadable. Callers that omit ``round`` opt out of dedup on purpose — they + have no round identity to key on. + """ + + def setup_method(self): + import tempfile + self.tmpdir = tempfile.mkdtemp() + + def test_same_round_same_message_recorded_once(self): + log = SessionLog(self.tmpdir, session_key="err_dedup") + for _ in range(5): + log.record_error({ + "message": "APIError: <400> boom", + "recoverable": False, + "round": 15, + }) + assert len(log.get_errors()) == 1 + + def test_round_and_message_changes_are_kept(self): + log = SessionLog(self.tmpdir, session_key="err_kept") + log.record_error({"message": "APIError: boom", "round": 15}) + log.record_error({"message": "APIError: boom", "round": 16}) + log.record_error({"message": "different failure", "round": 15}) + assert len(log.get_errors()) == 3 + + def test_callers_without_a_round_opt_out_of_dedup(self): + log = SessionLog(self.tmpdir, session_key="err_optout") + log.record_error({"message": "no round identity"}) + log.record_error({"message": "no round identity"}) + assert len(log.get_errors()) == 2 From 4b5f976039b88938b9f2162272dbea88065825ce Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Wed, 5 Aug 2026 01:16:24 +0800 Subject: [PATCH 06/14] Use native Windows shell semantics in the local code executor --- ms_agent/tools/code/local_code_executor.py | 19 +++++++- .../tools/test_local_code_executor_windows.py | 47 +++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) create mode 100644 tests/tools/test_local_code_executor_windows.py diff --git a/ms_agent/tools/code/local_code_executor.py b/ms_agent/tools/code/local_code_executor.py index 261859278..e46a6911f 100644 --- a/ms_agent/tools/code/local_code_executor.py +++ b/ms_agent/tools/code/local_code_executor.py @@ -332,6 +332,18 @@ def _build_env(self, field: str, inherit: bool = False) -> Dict[str, str]: 'HOME': os.environ.get('HOME', ''), 'LANG': os.environ.get('LANG', ''), } + if os.name == 'nt': + # ``create_subprocess_shell`` uses the native Windows command + # processor. Keep the non-secret OS/user variables that cmd + # and programs using temporary/profile directories require. + for key in ( + 'SYSTEMROOT', 'WINDIR', 'COMSPEC', 'PATHEXT', 'TEMP', + 'TMP', 'TMPDIR', 'USERPROFILE', 'HOMEDRIVE', + 'HOMEPATH', 'USERNAME', 'APPDATA', 'LOCALAPPDATA', + 'PROGRAMDATA', 'OS', 'PROCESSOR_ARCHITECTURE'): + value = os.environ.get(key) + if value is not None: + env[key] = value if not self.tool_config or not hasattr(self.tool_config, field): return env @@ -589,7 +601,12 @@ async def call_tool(self, server_name: str, *, tool_name: str, indent=2) def _prepare_shell_command(self, command: str) -> str: - """Wrap composite shell input in ``sh -lc`` when needed (matches sandbox behavior).""" + """Use the native Windows shell; wrap composite POSIX input when needed.""" + if os.name == 'nt': + # asyncio.create_subprocess_shell delegates to cmd.exe on Windows; + # wrapping native syntax in a usually unavailable POSIX shell + # makes otherwise valid compound commands fail. + return command shell_meta = ('&&', '||', '|', ';', '>', '<', '`', '$(', 'cd ', 'export ') already_wrapped = command.lstrip().startswith( diff --git a/tests/tools/test_local_code_executor_windows.py b/tests/tools/test_local_code_executor_windows.py new file mode 100644 index 000000000..8ba819c6e --- /dev/null +++ b/tests/tools/test_local_code_executor_windows.py @@ -0,0 +1,47 @@ +import os +from unittest import mock + +from ms_agent.tools.code.local_code_executor import LocalCodeExecutionTool + + +def _bare_tool() -> LocalCodeExecutionTool: + """Build a unit-test instance without starting kernels or checking deps.""" + tool = LocalCodeExecutionTool.__new__(LocalCodeExecutionTool) + tool.tool_config = None + return tool + + +def test_composite_command_uses_native_windows_shell(): + command = 'cd work && echo ok > result.txt' + with mock.patch('ms_agent.tools.code.local_code_executor.os.name', 'nt'): + assert _bare_tool()._prepare_shell_command(command) == command + + +def test_sanitized_env_keeps_windows_runtime_variables(): + windows_env = { + 'PATH': r'C:\Windows\System32', + 'SYSTEMROOT': r'C:\Windows', + 'WINDIR': r'C:\Windows', + 'COMSPEC': r'C:\Windows\System32\cmd.exe', + 'PATHEXT': '.COM;.EXE;.BAT;.CMD', + 'TEMP': r'C:\Users\tester\AppData\Local\Temp', + 'TMP': r'C:\Users\tester\AppData\Local\Temp', + 'USERPROFILE': r'C:\Users\tester', + 'HOMEDRIVE': 'C:', + 'HOMEPATH': r'\Users\tester', + 'USERNAME': 'tester', + 'APPDATA': r'C:\Users\tester\AppData\Roaming', + 'LOCALAPPDATA': r'C:\Users\tester\AppData\Local', + 'SECRET_TOKEN': 'must-not-leak', + } + with mock.patch.dict(os.environ, windows_env, clear=True), mock.patch( + 'ms_agent.tools.code.local_code_executor.os.name', 'nt'): + env = _bare_tool()._build_env('shell_env', inherit=False) + + for key in ( + 'SYSTEMROOT', 'WINDIR', 'COMSPEC', 'PATHEXT', 'TEMP', 'TMP', + 'USERPROFILE', 'HOMEDRIVE', 'HOMEPATH', 'USERNAME', 'APPDATA', + 'LOCALAPPDATA'): + assert env[key] == windows_env[key] + assert env['INHERITED_FROM_LOCAL'] == 'False' + assert 'SECRET_TOKEN' not in env From 0ae72247dc244601d6e2d7586a43428935bbd1cf Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Thu, 6 Aug 2026 16:35:56 +0800 Subject: [PATCH 07/14] Infer provider from model name only when no service is configured --- ms_agent/llm/router.py | 16 ++++++++++---- tests/llm/test_provider_layer.py | 38 ++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/ms_agent/llm/router.py b/ms_agent/llm/router.py index 1195fe1f4..5e46659cc 100644 --- a/ms_agent/llm/router.py +++ b/ms_agent/llm/router.py @@ -3,9 +3,9 @@ ``ProviderRouter.create(config)`` is the data-driven replacement for the hard-coded ``all_services_mapping`` factory. It resolves the spec (by service -name, then by model-name keywords, else a generic OpenAI-compatible fallback), -resolves credentials, builds the matching transport, and returns an -``LLMProvider``. +name; by model-name keywords only when no service is configured; else a generic +OpenAI-compatible fallback named after the service), resolves credentials, +builds the matching transport, and returns an ``LLMProvider``. ``LLMProvider`` is a drop-in for the legacy LLM instances on the agent hot path: it exposes ``.model``, ``.config`` and ``.generate(messages, tools, **kwargs)`` @@ -108,7 +108,15 @@ def create(self, config: DictConfig) -> LLMProvider: model = config.llm.model spec = self._registry.get(service) - if spec is None: + if spec is None and not service: + # Model-name inference only applies when the config names no + # service. A configured-but-unknown service is a custom provider: + # its credentials live under ``_api_key`` / + # ``_base_url``, so inferring a built-in spec from the + # model name would look them up under that vendor's name instead + # and silently fall back to that vendor's default endpoint (e.g. a + # private gateway serving a model named ``deepseek-*`` would be + # routed to api.deepseek.com). spec = self._registry.resolve_by_model(model) if spec is None: logger.info( diff --git a/tests/llm/test_provider_layer.py b/tests/llm/test_provider_layer.py index 9e05fd481..fd26003d3 100644 --- a/tests/llm/test_provider_layer.py +++ b/tests/llm/test_provider_layer.py @@ -146,6 +146,44 @@ def test_explicit_true_forces_router(self): self.assertIsInstance(obj, LLMProvider) +class TestCustomServiceRouting(unittest.TestCase): + """A configured-but-unknown service names a custom provider: it must not be + re-resolved to a built-in spec via the model name, or its credentials and + endpoint (stored under ``_*``) would be looked up under the wrong + provider name.""" + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_custom_service_keeps_its_name_and_endpoint(self): + from ms_agent.llm.router import ProviderRouter + config = OmegaConf.create({ + 'llm': { + 'service': 'ms-test', + # Vendor keyword in the model name must not hijack the spec. + 'model': 'deepseek-v4-flash', + 'ms-test_api_key': 'sk-custom', + 'ms-test_base_url': 'https://gateway.invalid/compatible-mode/v1', + } + }) + provider = ProviderRouter().create(config) + self.assertEqual('ms-test', provider.spec.name) + self.assertEqual( + 'sk-custom', + CredentialResolver.resolve_api_key(provider.spec, config)) + self.assertEqual( + 'https://gateway.invalid/compatible-mode/v1', + CredentialResolver.resolve_base_url(provider.spec, config)) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_model_inference_still_applies_without_service(self): + from ms_agent.llm.router import ProviderRouter + config = OmegaConf.create( + {'llm': { + 'model': 'deepseek-v4-flash', + 'deepseek_api_key': 'sk-x' + }}) + self.assertEqual('deepseek', ProviderRouter().create(config).spec.name) + + class TestCredentialResolver(unittest.TestCase): @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') From 738640e8941cd2f0954ba2de195a29c04c5becf3 Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Mon, 10 Aug 2026 16:48:30 +0800 Subject: [PATCH 08/14] Ingest a round's memory before the blocking interactive input wait --- .gitignore | 2 ++ ms_agent/agent/llm_agent.py | 15 ++++++++++++--- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/.gitignore b/.gitignore index 534cc302a..0f46edaa3 100644 --- a/.gitignore +++ b/.gitignore @@ -171,3 +171,5 @@ webui/work_dir/ .ms_agent_snapshots/ .ms_agent/ + +webui/frontend/.react-router \ No newline at end of file diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index 03d30a0fb..0276944be 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -2248,6 +2248,18 @@ async def run_loop(self, messages: Union[List[Message], str], for msg in messages[pre_step_len:step_end_len]: self.session_log.append(self._msg_to_dict(msg)) + # Ingest THIS round's memory here, for the same reason the + # SessionLog write above moved up: after_tool_call blocks on the + # next prompt in interactive mode. Running afterwards made + # 'add_after_step' mean "after the *next* step" — round N was + # only ingested once round N+1 arrived, a single-round session + # was never ingested at all (the run-loop tail that would have + # caught it is skipped when the input source raises EOF), and + # the message list it saw had the next user turn already + # appended. Before the block, `messages` is exactly this round. + await self.add_memory( + messages, add_type='add_after_step', **kwargs) + await self.after_tool_call(messages) self.runtime.round += 1 @@ -2258,9 +2270,6 @@ async def run_loop(self, messages: Union[List[Message], str], self.session_log.append(self._msg_to_dict(msg)) self.session_log.round = self.runtime.round - # save memory and history - await self.add_memory( - messages, add_type='add_after_step', **kwargs) self.save_history(messages) # +1 means the next round the assistant may give a conclusion From ed0175326e07312cb20ce7c8da99e968ffd31c8b Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Mon, 10 Aug 2026 16:48:31 +0800 Subject: [PATCH 09/14] Release the mem0 vector client on close and log ingestion failures --- .../memory/unified/backends/mem0_adapter.py | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/ms_agent/memory/unified/backends/mem0_adapter.py b/ms_agent/memory/unified/backends/mem0_adapter.py index 1798ea044..ffdc2c681 100644 --- a/ms_agent/memory/unified/backends/mem0_adapter.py +++ b/ms_agent/memory/unified/backends/mem0_adapter.py @@ -84,6 +84,18 @@ async def start(self, **kwargs: Any) -> None: self._mem0 = None async def close(self) -> None: + # Drop the vector client explicitly. Embedded stores (qdrant/chroma on a + # local path) hold an exclusive OS file lock, so merely releasing the + # reference leaves the store locked until GC gets around to it -- long + # enough that the next agent, or any other process on the same path, + # fails with "already accessed by another instance". + client = getattr(getattr(self._mem0, 'vector_store', None), 'client', + None) + if client is not None: + try: + client.close() + except Exception as e: # pragma: no cover - best-effort teardown + logger.debug(f'[mem0_backend] vector client close failed: {e}') self._mem0 = None # ── inject ─────────────────────────────────────────────────────── @@ -142,7 +154,14 @@ async def on_messages( return await _offload(self._mem0.add, convo, user_id=self._user_id) except Exception as e: - logger.warning(f'[mem0_backend] add failed: {e}') + # Deliberately non-fatal -- a memory write must never break the + # turn. But log at error with the exception type: this is the only + # trace a failed ingestion leaves, and the symptom it produces + # (conversation succeeds, memory stays empty forever) gives the + # user nothing to search for. + logger.error( + f'[mem0_backend] add failed, nothing was persisted for this ' + f'round: {type(e).__name__}: {e}') # ── Search ─────────────────────────────────────────────────────── From 4335a8b981f53693d21f797b02f93f370775a920 Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Mon, 10 Aug 2026 21:57:16 +0800 Subject: [PATCH 10/14] Let a handler that supports it receive a round's parallel permission asks at once --- ms_agent/permission/enforcer.py | 48 +++++++++++--- ms_agent/permission/handler.py | 20 ++++++ tests/permission/test_parallel_permission.py | 67 ++++++++++++++++++-- 3 files changed, 122 insertions(+), 13 deletions(-) diff --git a/ms_agent/permission/enforcer.py b/ms_agent/permission/enforcer.py index 4cc44aff9..ca10b020b 100644 --- a/ms_agent/permission/enforcer.py +++ b/ms_agent/permission/enforcer.py @@ -40,10 +40,13 @@ def __init__( self._memory = memory or PermissionMemory() self._matcher = PermissionMatcher() # Parallel tool calls (asyncio.gather in ToolManager.parallel_call_tool) - # would otherwise invoke the interactive handler concurrently — N - # prompts fighting over one terminal deadlocks. Serialize asks with a - # lock created lazily per running loop (the per-turn TUI uses a fresh - # loop each turn, so a single init-time Lock would bind to the wrong one). + # reach the handler concurrently. Whether that is safe is the HANDLER's + # property, not a blanket rule: a terminal-bound one (CLI prompt / TUI + # menu) deadlocks with N prompts fighting over one stdin, while a + # request_id-keyed UI wants them all at once. Handlers opt in with + # ``supports_concurrent_asks``; everyone else is serialized with a lock + # created lazily per running loop (the per-turn TUI uses a fresh loop + # each turn, so a single init-time Lock would bind to the wrong one). self._ask_lock: asyncio.Lock | None = None self._ask_lock_loop = None @@ -54,13 +57,31 @@ def _ask_lock_for_loop(self) -> 'asyncio.Lock': self._ask_lock_loop = loop return self._ask_lock - async def _serialized_ask(self, **kwargs) -> PermissionResponse: + async def _ask_user(self, + *, + forced: bool = False, + **kwargs) -> PermissionResponse | None: + """Put one ask in front of the user, serialized unless the handler + declares it can service several at once. + + Returns ``None`` when a queued ask turned out to be unnecessary: while + it waited for the lock, an earlier ask in the same round was answered + with allow_session / allow_always covering this call too, so prompting + again would ask the user something they just answered. ``forced`` asks + (a SafetyGuard confirmation) skip that shortcut — memory must never + bypass a safety ask. + """ # ``call_id`` is a newer, optional kwarg (see check()). A handler that # predates it — or a lightweight test double — need not accept it; drop # it for such handlers so their fixed signature keeps working. if 'call_id' in kwargs and not self._handler_accepts('call_id'): kwargs.pop('call_id') + if getattr(self._handler, 'supports_concurrent_asks', False): + return await self._handler.ask(**kwargs) async with self._ask_lock_for_loop(): + if not forced and self._memory.matches(kwargs['tool_name'], + kwargs['tool_args']): + return None return await self._handler.ask(**kwargs) def _handler_accepts(self, param: str) -> bool: @@ -96,7 +117,8 @@ async def check( if force_decision and force_decision.action == 'ask': suggestions = generate_suggestions(tool_name, tool_args) - response = await self._serialized_ask( + response = await self._ask_user( + forced=True, tool_name=tool_name, tool_args=tool_args, context=force_decision.reason or '', @@ -126,9 +148,9 @@ async def check( reason='Allowed by remembered permission', ) - # 5. Ask user via handler (serialized against parallel tool calls) + # 5. Ask user via handler (serialized unless it opts into concurrency) suggestions = generate_suggestions(tool_name, tool_args) - response = await self._serialized_ask( + response = await self._ask_user( tool_name=tool_name, tool_args=tool_args, context='', @@ -140,10 +162,18 @@ async def check( def _process_response( self, - response: PermissionResponse, + response: PermissionResponse | None, tool_name: str, tool_args: dict[str, Any], ) -> PermissionDecision: + if response is None: + # The ask was skipped: memory started covering this call while it + # was queued behind another one (see _ask_user). + return PermissionDecision( + action='allow', + reason='Allowed by remembered permission', + ) + if response.action == PermissionAction.ALLOW_ONCE: return PermissionDecision( action='allow', reason='User allowed once') diff --git a/ms_agent/permission/handler.py b/ms_agent/permission/handler.py index a12590179..9e9802c57 100644 --- a/ms_agent/permission/handler.py +++ b/ms_agent/permission/handler.py @@ -34,6 +34,17 @@ class PermissionResponse: class PermissionHandler(Protocol): + """Confirmation UI for a tool call the policy can't decide on its own. + + Optional duck-typed attribute ``supports_concurrent_asks`` (default + ``False`` when absent) declares whether several asks may be in flight at + once. It is False for anything bound to the one terminal — N prompts + fighting over a single stdin/menu deadlock — so ``PermissionEnforcer`` + serializes those. A handler that keys pending asks by id and renders them + independently (``WebPermissionHandler``) sets it True, so a round's + parallel tool calls all surface for decision at the same time instead of + one-at-a-time behind whoever the user answers first. + """ async def ask( self, @@ -49,6 +60,9 @@ async def ask( class AutoPermissionHandler: """Always allows — used as fallback or in auto mode.""" + # Never blocks on anything, so it has no reason to be serialized. + supports_concurrent_asks = True + async def ask( self, tool_name: str, @@ -141,6 +155,12 @@ def emit(self, event: dict[str, Any]) -> None: class WebPermissionHandler: """Async handler that suspends on a Future until the frontend responds.""" + # Pending asks are keyed by request_id and each renders as its own card, so + # a round's parallel tool calls can all wait for a decision simultaneously. + # Serializing them instead would show one card at a time while the untouched + # siblings sat there looking like they were already running. + supports_concurrent_asks = True + def __init__( self, event_emitter: EventEmitter, diff --git a/tests/permission/test_parallel_permission.py b/tests/permission/test_parallel_permission.py index 4f48f3f78..add905087 100644 --- a/tests/permission/test_parallel_permission.py +++ b/tests/permission/test_parallel_permission.py @@ -1,7 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Parallel tool calls (ToolManager.parallel_call_tool → asyncio.gather) must -not invoke the interactive permission handler concurrently: N prompts fighting -over one terminal deadlocks. The enforcer serializes asks.""" +"""Parallel tool calls (ToolManager.parallel_call_tool → asyncio.gather) reach +the permission handler concurrently. Whether that is safe is the handler's +property: a terminal-bound one deadlocks with N prompts fighting over one +stdin, so the enforcer serializes it; a handler that declares +``supports_concurrent_asks`` gets them all at once (a web UI renders each +pending ask as its own card).""" import asyncio import pytest @@ -45,9 +48,65 @@ async def test_parallel_asks_are_serialized(tmp_path): assert handler.max_in_flight == 1 +class _ConcurrentProbe(_ConcurrencyProbe): + """Same probe, but declaring it can service several asks at once.""" + supports_concurrent_asks = True + + +@pytest.mark.asyncio +async def test_concurrent_handler_asks_overlap(tmp_path): + """A handler that opts in sees the whole round's asks at once, instead of + one-at-a-time behind whichever the user answers first.""" + cfg = PermissionConfig.from_dict({'mode': 'restricted'}) + handler = _ConcurrentProbe() + enf = PermissionEnforcer( + config=cfg, handler=handler, + memory=PermissionMemory(project_path=str(tmp_path))) + + results = await asyncio.gather( + *[enf.check('some_tool', {'i': i}) for i in range(5)]) + + assert all(r.action == 'allow' for r in results) + assert handler.calls == 5 + assert handler.max_in_flight == 5 + + +class _AlwaysAllowOnce: + """Serialized handler whose FIRST answer is allow_always; later asks should + never reach it — memory now covers them.""" + + def __init__(self): + self.calls = 0 + + async def ask(self, tool_name, tool_args, context, suggestions=None, + call_id=''): + self.calls += 1 + await asyncio.sleep(0.01) + return PermissionResponse( + action=PermissionAction.ALLOW_ALWAYS, pattern=tool_name) + + +@pytest.mark.asyncio +async def test_queued_ask_skipped_once_memory_covers_it(tmp_path): + """A serialized ask waiting its turn re-checks memory before prompting: the + user already answered "always allow" for this exact pattern on the sibling + ahead of it in the queue.""" + cfg = PermissionConfig.from_dict({'mode': 'restricted'}) + handler = _AlwaysAllowOnce() + enf = PermissionEnforcer( + config=cfg, handler=handler, + memory=PermissionMemory(project_path=str(tmp_path))) + + results = await asyncio.gather( + *[enf.check('some_tool', {'i': i}) for i in range(4)]) + + assert all(r.action == 'allow' for r in results) + assert handler.calls == 1 # the other three were covered by memory + + class _CallIdCapture: """Handler that records the call_id it was asked with (and tolerates - handlers that don't accept it — see enforcer._serialized_ask).""" + handlers that don't accept it — see enforcer._ask_user).""" def __init__(self): self.seen = [] From c1654a57fb4bc4eda7f634c849b45814af9b95a9 Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Mon, 10 Aug 2026 21:57:34 +0800 Subject: [PATCH 11/14] Report each parallel tool call's completion as it finishes, not after the batch --- ms_agent/agent/llm_agent.py | 91 ++++++++++++++++------------ ms_agent/tools/tool_manager.py | 30 ++++++++- ms_agent/tui/app.py | 8 ++- ms_agent/tui/permission.py | 34 +++++++++-- ms_agent/tui/renderer.py | 64 ++++++++++++++++--- tests/tui/test_permission_handler.py | 66 ++++++++++++++++++++ tests/tui/test_renderer.py | 39 ++++++++++++ 7 files changed, 275 insertions(+), 57 deletions(-) create mode 100644 tests/tui/test_permission_handler.py diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index d9f0d1a7e..1efe70886 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -678,23 +678,29 @@ async def parallel_tool_call(self, """ Execute multiple tool calls in parallel and append results to the message list. + Each result is turned into its tool message and ANNOUNCED + (``ToolCallCompleted``, plus ``PlanUpdated`` for todo tools) the moment + that call returns — not once the whole batch has. A round can contain + one call blocked on a human's approval next to calls that are already + finished, and those must be able to report themselves as finished. + The messages are still appended in call order, matching ``tool_calls``. + Args: messages (List[Message]): Current conversation history. Returns: List[Message]: Updated message list including tool responses. """ - tool_call_result = await self.tool_manager.parallel_call_tool( - messages[-1].tool_calls) - assert len(tool_call_result) == len(messages[-1].tool_calls) - for tool_call_result, tool_call_query in zip(tool_call_result, - messages[-1].tool_calls): - tool_call_result_format = ToolResult.from_raw(tool_call_result) + tool_calls = messages[-1].tool_calls + results: Dict[int, Message] = {} + + def _on_result(index: int, tool_call, raw, duration_s: float) -> None: + tool_call_result_format = ToolResult.from_raw(raw) _new_message = Message( role='tool', content=tool_call_result_format.text, - tool_call_id=tool_call_query['id'], - name=tool_call_query['tool_name'], + tool_call_id=tool_call['id'], + name=tool_call['tool_name'], resources=tool_call_result_format.resources, tool_detail=tool_call_result_format.tool_detail, hook_attachments=tool_call_result_format.hook_attachments, @@ -704,11 +710,42 @@ async def parallel_tool_call(self, if _new_message.tool_call_id is None: # If tool call id is None, add a random one _new_message.tool_call_id = str(uuid.uuid4())[:8] - tool_call_query['id'] = _new_message.tool_call_id - messages.append(_new_message) + tool_call['id'] = _new_message.tool_call_id + # This call's OWN wall clock. It used to be the batch's span + # attributed to every call in it, which over-reported every fast + # tool that shared a round with a slow one. + _new_message._duration_ms = int(duration_s * 1000) + results[index] = _new_message self.log_output(_new_message.content) + self._emit_tool_completed(_new_message, duration_s) + + await self.tool_manager.parallel_call_tool( + tool_calls, on_result=_on_result) + assert len(results) == len(tool_calls) + for index in range(len(tool_calls)): + messages.append(results[index]) return messages + def _emit_tool_completed(self, message: Message, duration_s: float) -> None: + """Report one finished tool call to the UI (and the plan it may carry).""" + if self._event_sink is None: + return + content = ( + message.content + if isinstance(message.content, str) else str(message.content)) + is_error = bool(getattr(message, 'is_error', False)) + self._event_sink.emit( + ToolCallCompleted( + call_id=str(getattr(message, 'tool_call_id', '') or ''), + name=str(getattr(message, 'name', '') or ''), + result=content or '', + error=(content or 'tool call failed') if is_error else None, + duration_s=round(duration_s, 3))) + # todo/split_task tool results drive the plan panel. + plan = self._extract_plan_from_tool_result(message) + if plan is not None: + self._event_sink.emit(PlanUpdated(entries=plan)) + def _select_permission_handler(self, mode: str): """Pick the PermissionHandler by mode + runtime environment. @@ -1727,37 +1764,11 @@ def _next_chunk(_g=_gen): name=str( tc.get('tool_name') or tc.get('name') or ''), arguments=tc.get('arguments'))) - _tool_start = len(messages) - _tool_t0 = time.monotonic() + # Each call stamps its own ``_duration_ms`` (persistence/replay) and + # emits its own ToolCallCompleted / PlanUpdated as it finishes — + # inside parallel_tool_call, not after the batch, so one call held + # up by an approval doesn't hold back its finished siblings. messages = await self.parallel_tool_call(messages) - # Batch wall-clock: exact for the common single-tool round; for - # parallel multi-tool rounds it attributes the batch span to each - # (they ran concurrently within it). Stamp for persistence/replay - # regardless of sink, and report it live via ToolCallCompleted. - _tool_ms = int((time.monotonic() - _tool_t0) * 1000) - for m in messages[_tool_start:]: - if getattr(m, 'role', None) == 'tool': - m._duration_ms = _tool_ms - if self._event_sink is not None: - for m in messages[_tool_start:]: - if getattr(m, 'role', None) == 'tool': - _content = ( - m.content - if isinstance(m.content, str) else str(m.content)) - _is_err = bool(getattr(m, 'is_error', False)) - self._event_sink.emit( - ToolCallCompleted( - call_id=str( - getattr(m, 'tool_call_id', '') or ''), - name=str(getattr(m, 'name', '') or ''), - result=_content or '', - error=(_content or 'tool call failed') - if _is_err else None, - duration_s=round(_tool_ms / 1000, 3))) - # todo/split_task tool results drive the plan panel. - _plan = self._extract_plan_from_tool_result(m) - if _plan is not None: - self._event_sink.emit(PlanUpdated(entries=_plan)) # usage # NOTE: token accounting must run BEFORE after_tool_call. The interactive diff --git a/ms_agent/tools/tool_manager.py b/ms_agent/tools/tool_manager.py index 5a4074ff4..dac7dcf6e 100644 --- a/ms_agent/tools/tool_manager.py +++ b/ms_agent/tools/tool_manager.py @@ -8,6 +8,7 @@ import math import os import sys +import time import uuid from copy import copy from types import TracebackType @@ -744,8 +745,33 @@ async def single_call_tool(self, tool_info: ToolCall): True, } - async def parallel_call_tool(self, tool_list: List[ToolCall]): - tasks = [self.single_call_tool(tool) for tool in tool_list] + async def parallel_call_tool( + self, + tool_list: List[ToolCall], + on_result: Optional[Callable[[int, ToolCall, Any, float], + None]] = None, + ): + """Run a round's tool calls concurrently, in call order in the result. + + ``on_result(index, tool_call, result, duration_s)`` — when given — fires + the moment THAT call returns, rather than after the whole batch. The + distinction matters as soon as one call can block for a long time: under + interactive permissions a call suspended on a human's approval used to + hold back every sibling's completion, so calls that were already done + (or needed no approval at all) still looked like they were running until + the human answered. It runs on the event loop between tool calls, so + keep it cheap and non-throwing — an exception propagates out of the + gather and fails the round. + """ + + async def _call(index: int, tool: ToolCall): + started = time.monotonic() + result = await self.single_call_tool(tool) + if on_result is not None: + on_result(index, tool, result, time.monotonic() - started) + return result + + tasks = [_call(i, tool) for i, tool in enumerate(tool_list)] result = await asyncio.gather(*tasks) return result diff --git a/ms_agent/tui/app.py b/ms_agent/tui/app.py index 9c5ceb83d..aa25d4800 100644 --- a/ms_agent/tui/app.py +++ b/ms_agent/tui/app.py @@ -131,7 +131,13 @@ def __init__( # interactive at runtime and get real confirmations. from ms_agent.tui.permission import TUIPermissionHandler self.agent.set_permission_handler( - TUIPermissionHandler(console=self.console, theme=self.theme)) + TUIPermissionHandler( + console=self.console, + theme=self.theme, + # Lets the menu hold the renderer's draws while it owns the + # terminal (a sibling tool finishing mid-menu must not print + # into it) — see RichEventSink.hold_output. + renderer=self.renderer)) # ('new', None) | ('resume', '<#|id>') | None, set by session commands. self._pending_switch: Optional[Tuple[str, Optional[str]]] = None diff --git a/ms_agent/tui/permission.py b/ms_agent/tui/permission.py index 3a0f1ee3a..a031199cc 100644 --- a/ms_agent/tui/permission.py +++ b/ms_agent/tui/permission.py @@ -3,14 +3,20 @@ Matches the pattern used by Claude Code / Qoder / hermes: a compact header for the tool call, then a selectable menu (``❯`` cursor, ↑/↓ + number keys, Enter) -instead of a "type a letter" prompt. The enforcer serializes asks and this runs -on the main event loop, so a prompt_toolkit menu composes without terminal -contention; a non-TTY fallback keeps it scriptable. +instead of a "type a letter" prompt. A non-TTY fallback keeps it scriptable. + +This handler does NOT declare ``supports_concurrent_asks``, so the enforcer +serializes its asks — one terminal, one menu at a time. That alone is no longer +enough for a quiet screen: the menu runs on the same event loop the renderer +draws from, and a sibling tool call approved a moment earlier can finish while +this menu is up. So ``ask()`` also holds the renderer's output for its duration +(``RichEventSink.hold_output``). """ from __future__ import annotations import asyncio import json +from contextlib import contextmanager from rich.console import Console from typing import Any, Optional @@ -26,11 +32,31 @@ class TUIPermissionHandler: def __init__(self, console: Optional[Console] = None, io: Any = None, - theme: Theme = DEFAULT_THEME) -> None: + theme: Theme = DEFAULT_THEME, + renderer: Any = None) -> None: self._console = console or Console() self._theme = theme + # The event renderer, so its draws can be held while this menu owns the + # terminal. A sibling tool call finishing mid-menu would otherwise print + # its result line straight through the prompt_toolkit app the user is + # reading (tool completions arrive per call now, so that overlap is + # reachable whenever one call is approved while another is still asked). + self._renderer = renderer + + @contextmanager + def _own_screen(self): + hold = getattr(self._renderer, 'hold_output', None) + if hold is None: + yield # no renderer wired (tests, embedders) — nothing to hold + return + with hold(): + yield async def ask(self, tool_name, tool_args, context, suggestions=None): + with self._own_screen(): + return await self._ask(tool_name, tool_args, context, suggestions) + + async def _ask(self, tool_name, tool_args, context, suggestions=None): from ms_agent.permission.handler import (PermissionAction, PermissionResponse) suggestion = suggestions[0] if suggestions else tool_name diff --git a/ms_agent/tui/renderer.py b/ms_agent/tui/renderer.py index cb2268c08..d491b6550 100644 --- a/ms_agent/tui/renderer.py +++ b/ms_agent/tui/renderer.py @@ -15,6 +15,7 @@ import json import re import time +from contextlib import contextmanager from rich.console import Console from rich.markdown import Markdown from rich.markup import escape @@ -71,6 +72,9 @@ def __init__(self, # add a blank line between sections (tools ↔ assistant) for breathing # room without spacing tightly-grouped tool lines apart. self._last_kind: Optional[str] = None + # Buffered draws while something else owns the screen (see hold_output). + # None = draw straight through. + self._held: Optional[list] = None # ── sink protocol ────────────────────────────────────────────────────── @@ -79,6 +83,44 @@ def emit(self, event: AgentEvent) -> None: if handler is not None: handler(event) + # ── screen ownership ─────────────────────────────────────────────────── + + @contextmanager + def hold_output(self): + """Buffer this sink's draws while a modal owns the terminal. + + A permission menu (``tui.select.select_async``) is a prompt_toolkit + Application rendered inline on the SAME event loop this sink draws from, + so any print landing mid-menu corrupts what the user is reading. Tool + completions now arrive as each call finishes rather than after the whole + round (``LLMAgent.parallel_tool_call``), which is exactly when that can + happen: approving one call lets the next one's menu open while the first + is still executing, and its result lands underneath. + + Events are still HANDLED immediately — only the drawing waits — so + durations and internal state stay measured at the true moment. + + Nested holds share the outermost buffer; the flush happens once, when + the screen is actually released. + """ + if self._held is not None: + yield # already held by an outer scope — it owns the flush + return + self._held = [] + try: + yield + finally: + held, self._held = self._held, None + for args, kwargs in held: + self.console.print(*args, **kwargs) + + def _print(self, *args, **kwargs) -> None: + """Draw now, or queue it if a modal currently owns the screen.""" + if self._held is not None: + self._held.append((args, kwargs)) + return + self.console.print(*args, **kwargs) + def finalize(self) -> None: """Tear down any in-flight Live region (call on error / turn abort).""" if self._live is not None: @@ -161,24 +203,26 @@ def _on_tool_call_started(self, ev) -> None: if self._last_kind != 'tool': # Blank before a new tool group (turn start or after text), but keep # consecutive tool lines tight together. - self.console.print() + self._print() self._tool_started[ev.call_id] = time.monotonic() header = escape(tool_header(ev.name, ev.arguments)) - self.console.print(f'[{self.theme.tool_bullet}]•[/] {header}') + self._print(f'[{self.theme.tool_bullet}]•[/] {header}') self._last_kind = 'tool' def _on_tool_call_completed(self, ev) -> None: - # Indented one-line summary: " └ 42 lines · 1.2s". + # Indented one-line summary: " └ 42 lines · 1.2s". Measured HERE even + # when the draw is held back for a permission menu — the elapsed time is + # the tool's, not the user's deliberation. start = self._tool_started.pop(ev.call_id, None) dur = f' · {time.monotonic() - start:.1f}s' if start else '' if ev.error: summary = escape(tool_summary(ev.result, ev.error)) - self.console.print( + self._print( f' [{self.theme.tool_error_border}]└[/] ' f'[{self.theme.tool_error_border}]{summary}[/][dim]{dur}[/]') else: summary = escape(tool_summary(ev.result)) - self.console.print(f' [dim]└ {summary}{dur}[/]') + self._print(f' [dim]└ {summary}{dur}[/]') # ── plan / notices / context / errors ────────────────────────────────── @@ -196,7 +240,7 @@ def _on_plan_updated(self, ev) -> None: e.get('content', '') if isinstance(e, dict) else getattr( e, 'content', '')) lines.append(f'{mark.get(status, "○")} {content}') - self.console.print( + self._print( Panel( '\n'.join(lines), title='plan', @@ -207,7 +251,7 @@ def _on_context_compacted(self, ev) -> None: detail = '' if ev.before_tokens and ev.after_tokens: detail = f' {ev.before_tokens}→{ev.after_tokens} tok' - self.console.print( + self._print( f'[{self.theme.notice_info}]· context compacted{detail} ·[/]') def _on_notice(self, ev) -> None: @@ -225,17 +269,17 @@ def _on_notice(self, ev) -> None: background_color='default') else: body = Text(ev.text) - self.console.print(Panel(body, border_style='dim', expand=False)) + self._print(Panel(body, border_style='dim', expand=False)) return style = { 'success': self.theme.notice_success, 'warning': self.theme.notice_warning }.get(ev.level, self.theme.notice_info) - self.console.print(f'[{style}]{ev.text}[/]') + self._print(f'[{style}]{ev.text}[/]') def _on_error(self, ev) -> None: self.finalize() - self.console.print( + self._print( Panel( f'[bold]{ev.message}[/]', title='error', diff --git a/tests/tui/test_permission_handler.py b/tests/tui/test_permission_handler.py new file mode 100644 index 000000000..7191f6c55 --- /dev/null +++ b/tests/tui/test_permission_handler.py @@ -0,0 +1,66 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""TUIPermissionHandler owns the terminal while its menu is up. + +The menu is a prompt_toolkit Application rendered inline on the SAME event loop +the renderer draws from, and tool completions now arrive as each call finishes +(LLMAgent.parallel_tool_call) rather than after the whole round — so approving +one call can open the next call's menu while the first is still executing, and +its result would land inside the menu the user is reading. +""" +from io import StringIO + +import pytest +from rich.console import Console + +from ms_agent.permission.handler import PermissionAction +from ms_agent.tui import permission as permission_mod +from ms_agent.tui.permission import TUIPermissionHandler +from ms_agent.tui.renderer import RichEventSink +from ms_agent.tui.state import TuiState +from ms_agent.ui.events import ToolCallCompleted, ToolCallStarted + + +def _renderer(): + console = Console(file=StringIO(), force_terminal=False, width=80) + return RichEventSink(console, TuiState()), console + + +@pytest.mark.asyncio +async def test_sibling_completion_does_not_print_into_the_menu(monkeypatch): + renderer, console = _renderer() + handler = TUIPermissionHandler(console=console, renderer=renderer) + renderer.emit( + ToolCallStarted( + call_id='c1', name='file_system---read_file', + arguments={'path': 'a.txt'})) + seen_during_menu = {} + + async def fake_menu(options, *, default=0, header=None): + # A sibling call finishes while the user is deciding. + renderer.emit( + ToolCallCompleted( + call_id='c1', name='file_system---read_file', result='aaa')) + seen_during_menu['out'] = console.file.getvalue() + return 0 # "Allow once" + + monkeypatch.setattr(permission_mod, 'select_async', fake_menu) + resp = await handler.ask('file_system---write_file', {'path': 'b.txt'}, '') + + assert resp.action == PermissionAction.ALLOW_ONCE + assert 'aaa' not in seen_during_menu['out'] # held while the menu was up + assert 'aaa' in console.file.getvalue() # drawn once the menu closed + + +@pytest.mark.asyncio +async def test_handler_without_a_renderer_still_works(monkeypatch): + """Embedders/tests may construct the handler bare — no renderer to hold.""" + console = Console(file=StringIO(), force_terminal=False, width=80) + handler = TUIPermissionHandler(console=console) + + async def fake_menu(options, *, default=0, header=None): + return 4 # "Deny" + + monkeypatch.setattr(permission_mod, 'select_async', fake_menu) + resp = await handler.ask('code_executor---shell', {'command': 'rm -rf /'}, + '') + assert resp.action == PermissionAction.DENY diff --git a/tests/tui/test_renderer.py b/tests/tui/test_renderer.py index 2413ee38d..80f6a3bd7 100644 --- a/tests/tui/test_renderer.py +++ b/tests/tui/test_renderer.py @@ -119,3 +119,42 @@ def test_unhandled_event_is_ignored(): def test_finalize_is_safe_without_live(): sink, _, _ = _sink() sink.finalize() # no active Live — must not raise + + +def test_hold_output_defers_draws_until_the_screen_is_released(): + """A permission menu owns the terminal while it is up; a sibling tool call + finishing mid-menu must not print into it. Tool completions arrive per call + now (LLMAgent.parallel_tool_call), so that overlap is reachable.""" + sink, console, _ = _sink() + sink.emit(ToolCallStarted(call_id='c1', name='file_system---read_file', + arguments={'path': 'a.txt'})) + with sink.hold_output(): + sink.emit(ToolCallCompleted(call_id='c1', + name='file_system---read_file', + result='aaa')) + sink.emit(PlanUpdated(entries=[PlanEntry('do X', 'completed')])) + held = _out(console) + assert 'aaa' not in held and 'do X' not in held + out = _out(console) + assert 'aaa' in out and 'do X' in out # flushed on release, in order + assert out.index('aaa') < out.index('do X') + + +def test_hold_output_flushes_even_if_the_menu_raises(): + sink, console, _ = _sink() + try: + with sink.hold_output(): + sink.emit(Notice(level='success', text='saved')) + raise KeyboardInterrupt # user hit Ctrl-C in the menu + except KeyboardInterrupt: + pass + assert 'saved' in _out(console) + + +def test_nested_holds_flush_once_at_the_outermost_release(): + sink, console, _ = _sink() + with sink.hold_output(): + with sink.hold_output(): + sink.emit(Notice(level='success', text='inner')) + assert 'inner' not in _out(console) # inner exit must NOT flush + assert 'inner' in _out(console) From 1100a380e977ef7d16bd2a664e2d02a191fda7c1 Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Tue, 11 Aug 2026 02:11:32 +0800 Subject: [PATCH 12/14] Take memory ingestion off the turn's critical path The orchestrator now owns the write discipline around a backend: - schedule_add() runs the extraction-LLM + embedding cost (seconds) in a background task; flush_pending() is the teardown barrier so the last write is never dropped, and an inline fallback keeps writes when no loop is running. - retrieval/ingestion/flush serialize on one per-store asyncio lock (embedded qdrant underneath is lock-free single-client code). - a content-hash delta ledger (/ingest_state.json) makes each ingest send only messages the store has not seen; hashes are recorded only after a confirmed write, so a failed ingest retries naturally. - ingest_status reports the last outcome (state/count/error/pending) so a UI can show memory working instead of silence. Mem0Backend: per-turn retrieval cache (rounds 2..N of a tool-calling turn reuse round 1's search instead of paying an embedding round-trip each), on_messages returns the event count and propagates failures -- the orchestrator is the swallow-and-report layer now and needs the exception to keep failed messages un-marked for retry. --- .../memory/unified/backends/mem0_adapter.py | 83 +++--- ms_agent/memory/unified/config.py | 7 +- ms_agent/memory/unified/orchestrator.py | 265 +++++++++++++++++- tests/memory/test_backend_contracts.py | 13 + tests/memory/test_orchestrator_scheduling.py | 224 +++++++++++++++ tests/memory/test_unified_memory.py | 10 +- 6 files changed, 551 insertions(+), 51 deletions(-) create mode 100644 tests/memory/test_orchestrator_scheduling.py diff --git a/ms_agent/memory/unified/backends/mem0_adapter.py b/ms_agent/memory/unified/backends/mem0_adapter.py index 43f8008c5..4323763aa 100644 --- a/ms_agent/memory/unified/backends/mem0_adapter.py +++ b/ms_agent/memory/unified/backends/mem0_adapter.py @@ -67,8 +67,13 @@ def __init__(self, config: MemoryConfig) -> None: self._config = config self._mem0: Any = None # mem0.Memory instance self._user_id: str = config.user_id - self._snapshot: Optional[str] = None - self._snapshot_dirty = True + # Per-turn retrieval cache: one turn = one embedding + one vector + # search. The turn key is the latest user message — every round of a + # multi-round (tool-calling) turn injects with the same user message, + # so rounds 2..N reuse the round-1 results instead of paying another + # embedding round-trip each. Invalidated on writes/deletes. + self._turn_cache_key: Optional[str] = None + self._turn_cache_results: Optional[list] = None # ── Lifecycle ──────────────────────────────────────────────────── @@ -111,13 +116,21 @@ async def inject( if not query: return messages - try: - results = _result_list(await _offload(_mem0_search, self._mem0, - query, self._user_id)) - if not results: + turn_key = f'{self._user_id}\x1f{query}' + if turn_key == self._turn_cache_key \ + and self._turn_cache_results is not None: + results = self._turn_cache_results + else: + try: + results = _result_list( + await _offload(_mem0_search, self._mem0, query, + self._user_id)) + except Exception as e: + logger.debug(f'[mem0_backend] search failed: {e}') return messages - except Exception as e: - logger.debug(f'[mem0_backend] search failed: {e}') + self._turn_cache_key = turn_key + self._turn_cache_results = results + if not results: return messages formatted = self._format_results(results) @@ -139,31 +152,32 @@ async def on_messages( self, messages: List[Dict[str, Any]], **kwargs: Any, - ) -> None: + ) -> int: + """Ingest via mem0's fact extraction. Returns the number of memory + events mem0 produced (ADD/UPDATE/DELETE). Raises on failure — the + orchestrator owns the swallow-and-report policy, and needs the + exception to know the write did NOT land (so its delta ledger keeps + the messages for a retry instead of marking them ingested).""" if not self._mem0: - return - try: - # mem0 rejects non-chat fields and roles like `tool`; feed it the - # user/assistant text turns only. - convo = [ - { - 'role': m['role'], - 'content': m['content'] - } for m in messages - if m.get('role') in ('user', 'assistant') and m.get('content') - ] - if not convo: - return - await _offload(self._mem0.add, convo, user_id=self._user_id) - except Exception as e: - # Deliberately non-fatal -- a memory write must never break the - # turn. But log at error with the exception type: this is the only - # trace a failed ingestion leaves, and the symptom it produces - # (conversation succeeds, memory stays empty forever) gives the - # user nothing to search for. - logger.error( - f'[mem0_backend] add failed, nothing was persisted for this ' - f'round: {type(e).__name__}: {e}') + return 0 + # mem0 rejects non-chat fields and roles like `tool`; feed it the + # user/assistant text turns only. + convo = [ + { + 'role': m['role'], + 'content': m['content'] + } for m in messages + if m.get('role') in ('user', 'assistant') and m.get('content') + ] + if not convo: + return 0 + result = await _offload(self._mem0.add, convo, user_id=self._user_id) + # A write changes what retrieval should see. + self._turn_cache_key = None + self._turn_cache_results = None + if isinstance(result, dict): + return len(result.get('results') or []) + return len(result or []) # ── Search ─────────────────────────────────────────────────────── @@ -191,8 +205,9 @@ async def search( # ── Cache ──────────────────────────────────────────────────────── def invalidate(self) -> None: - self._snapshot = None - self._snapshot_dirty = True + # External edit (UI delete, another writer): next inject re-queries. + self._turn_cache_key = None + self._turn_cache_results = None # ── Internal helpers ───────────────────────────────────────────── diff --git a/ms_agent/memory/unified/config.py b/ms_agent/memory/unified/config.py index 7fa247e5b..5b57cde2e 100644 --- a/ms_agent/memory/unified/config.py +++ b/ms_agent/memory/unified/config.py @@ -26,6 +26,11 @@ class MemoryConfig: # LLM for extraction (reuses agent LLM if None) llm_config: Optional[Dict[str, Any]] = None + # Ingest every Nth completed turn (1 = every turn). Turns skipped by the + # interval are still covered later: the orchestrator's delta ledger sends + # everything not yet ingested on the next firing ingest. + ingest_interval: int = 1 + # Backend-specific options keyed by backend name backend_options: Dict[str, Any] = field(default_factory=dict) @@ -113,7 +118,7 @@ def from_dict_config(cls, cfg: DictConfig) -> 'MemoryConfig': if k in ns: flat[k] = ns[k] - for k in ('enabled', 'base_dir', 'llm_config'): + for k in ('enabled', 'base_dir', 'llm_config', 'ingest_interval'): if k in raw: flat[k] = raw[k] diff --git a/ms_agent/memory/unified/orchestrator.py b/ms_agent/memory/unified/orchestrator.py index 789096a88..691787306 100644 --- a/ms_agent/memory/unified/orchestrator.py +++ b/ms_agent/memory/unified/orchestrator.py @@ -4,11 +4,39 @@ prompt injection, retrieval strategies, or tool definitions. All of that lives inside the MemoryBackend implementation selected by configuration. +What it DOES own is the write/read discipline around the backend: + +* **Serialization** — retrieval (``run``), ingestion (``add``) and flush all + take one asyncio lock per *store* (keyed by ``base_dir``), because embedded + vector stores (mem0 + local qdrant) have no internal locking at all and + mem0 even fans out worker threads inside ``add``. Lock per store, not per + orchestrator: ``SharedMemoryManager`` may hand different orchestrator + instances the same directory. +* **Background ingestion** — ``schedule_add`` takes the extraction-LLM + + embedding cost (seconds) off the turn's critical path. Tasks are retained + for ``flush_pending`` so a teardown cannot silently drop the last write. + If no loop is running the write happens inline — slower, never lost. +* **Delta ledger** — every ingested message's content hash is remembered + (in memory + ``/ingest_state.json``), so each ingest sends only + the messages the store has not seen. Without this, every ingest re-sent + the whole conversation: cost O(rounds x history), and re-extraction of old + turns. Hashes are only recorded AFTER a successful backend write, so a + failed ingest retries naturally on the next turn (the analog of a + watermark that only advances on a confirmed write). +* **Status** — ``ingest_status`` reports the last ingest outcome so a UI can + show "memory updated / failed" instead of silence. + Registered as ``unified_memory`` in ``memory_mapping``. """ from __future__ import annotations -from typing import Any, Dict, List, Optional +import asyncio +import hashlib +import json +import os +import tempfile +import time +from typing import Any, Dict, List, Optional, Set from ms_agent.llm.utils import Message from ms_agent.memory.base import Memory @@ -22,6 +50,32 @@ logger = get_logger() +# One lock per storage directory. Never per orchestrator instance: two +# orchestrators over the same path (possible through SharedMemoryManager's +# llm-dependent cache key) must still serialize against each other. +_STORE_LOCKS: Dict[str, asyncio.Lock] = {} + + +def _store_lock(base_dir: str) -> asyncio.Lock: + key = os.path.abspath(str(base_dir or '.')) + lock = _STORE_LOCKS.get(key) + if lock is None: + lock = _STORE_LOCKS.setdefault(key, asyncio.Lock()) + return lock + + +# Only conversational text is ingested (mirrors what backends extract from); +# system prompts and tool payloads never become long-term memory rows. +_INGEST_ROLES = ('user', 'assistant') + +_LEDGER_FILE = 'ingest_state.json' +_LEDGER_MAX = 4096 + + +def _content_hash(msg: Dict[str, Any]) -> str: + raw = f"{msg.get('role', '')}\x1f{msg.get('content', '')}" + return hashlib.sha256(raw.encode('utf-8')).hexdigest()[:16] + class MemoryOrchestrator(Memory): """Thin adapter between the ms-agent ``Memory`` ABC and a @@ -42,6 +96,17 @@ def __init__(self, config: Any) -> None: self.mem_config = self._parse_config(config) self._backend: Optional[MemoryBackend] = None self._started = False + # Background ingestion bookkeeping (see module docstring). + self._pending: Set[asyncio.Task] = set() + self._turns_since_ingest = 0 + self._ledger: Optional[Set[str]] = None # lazy-loaded from disk + self._ledger_order: List[str] = [] + self._status: Dict[str, Any] = { + 'state': 'idle', + 'at': None, + 'count': None, + 'error': None + } # ------------------------------------------------------------------ # Lazy backend construction @@ -71,21 +136,192 @@ async def run(self, messages: List[Message]) -> List[Message]: if not self.mem_config.enabled: return messages - backend = await self._ensure_started() - msg_dicts = _messages_to_dicts(messages) - injected = await backend.inject(msg_dicts) + # Retrieval must not overlap a write: the embedded stores underneath + # (qdrant local) are single-client, lock-free code. A scheduled ingest + # normally finishes while the user reads the previous answer, so this + # rarely actually waits. + async with _store_lock(self.mem_config.base_dir): + backend = await self._ensure_started() + msg_dicts = _messages_to_dicts(messages) + injected = await backend.inject(msg_dicts) return _dicts_to_messages(injected) # ------------------------------------------------------------------ - # Memory ABC -- add() + # Memory ABC -- add() / schedule_add() # ------------------------------------------------------------------ async def add(self, messages: List[Message], **kwargs: Any) -> None: - if not self.mem_config.enabled: + """Awaited ingestion (legacy callers / inline fallback).""" + if not self._should_ingest(): return - backend = await self._ensure_started() + await self._ingest(_messages_to_dicts(messages), **kwargs) + + def schedule_add(self, messages: List[Message], + **kwargs: Any) -> Optional[asyncio.Task]: + """Ingest in the background; returns the task (None when skipped or + run inline). The message list is snapshotted to dicts NOW — the + caller's list keeps mutating after this returns (the next user turn + is appended to it).""" + if not self._should_ingest(): + return None msg_dicts = _messages_to_dicts(messages) - await backend.on_messages(msg_dicts, **kwargs) + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is None: + # No loop to defer to — do the write inline rather than lose it. + asyncio.run(self._ingest(msg_dicts, **kwargs)) + return None + self._status.update(state='scheduled', error=None) + task = loop.create_task(self._ingest(msg_dicts, **kwargs)) + self._pending.add(task) + task.add_done_callback(self._pending.discard) + return task + + def _should_ingest(self) -> bool: + if not self.mem_config.enabled: + return False + interval = max(1, int(getattr(self.mem_config, 'ingest_interval', 1))) + self._turns_since_ingest += 1 + if self._turns_since_ingest < interval: + return False + self._turns_since_ingest = 0 + return True + + async def _ingest(self, msg_dicts: List[Dict[str, Any]], + **kwargs: Any) -> int: + """The single write path: locked, delta-only, status-reporting. + + Never raises — a memory write must not break anything above it; the + outcome lands in ``ingest_status`` instead. + """ + try: + async with _store_lock(self.mem_config.base_dir): + backend = await self._ensure_started() + delta = self._ledger_delta(msg_dicts) + if not delta: + self._set_status('ok', count=0) + return 0 + self._set_status('running') + result = await backend.on_messages(delta, **kwargs) + # Record hashes only after the backend accepted the write, so + # a failure leaves them un-marked and the next turn's delta + # carries them again (retry-by-construction). + self._ledger_mark(delta) + count = result if isinstance(result, int) else len(delta) + self._set_status('ok', count=count) + return count + except asyncio.CancelledError: + self._set_status('error', error='cancelled') + raise + except Exception as e: # noqa: BLE001 - reported via status + logger.error(f'[orchestrator] memory ingest failed, nothing was ' + f'persisted for this turn: {type(e).__name__}: {e}') + self._set_status('error', error=f'{type(e).__name__}: {e}') + return 0 + + def mark_ingested(self, messages: List[Message]) -> None: + """Advance the ledger WITHOUT ingesting (interrupted turns): a + half-finished answer must not be swept into the next turn's delta.""" + self._ledger_mark( + [ + m for m in _messages_to_dicts(messages) + if m.get('role') in _INGEST_ROLES and m.get('content') + ], + persist=True, + ) + + async def flush_pending(self, timeout: float = 15.0) -> None: + """Barrier: wait for scheduled ingests (teardown must not drop the + final write). Timeout guards against a wedged provider call.""" + pending = {t for t in self._pending if not t.done()} + if not pending: + return + done, still = await asyncio.wait(pending, timeout=timeout) + if still: + logger.warning( + f'[orchestrator] {len(still)} memory ingest(s) still running ' + f'after {timeout}s flush timeout') + + @property + def ingest_status(self) -> Dict[str, Any]: + status = dict(self._status) + status['pending'] = sum(1 for t in self._pending if not t.done()) + return status + + def _set_status(self, state: str, count: Optional[int] = None, + error: Optional[str] = None) -> None: + self._status.update( + state=state, + at=time.strftime('%Y-%m-%dT%H:%M:%S%z'), + count=count, + error=error) + + # ------------------------------------------------------------------ + # Ingest ledger (content hashes of already-ingested messages) + # ------------------------------------------------------------------ + + def _ledger_path(self) -> str: + return os.path.join(str(self.mem_config.base_dir), _LEDGER_FILE) + + def _ledger_load(self) -> Set[str]: + if self._ledger is None: + hashes: List[str] = [] + try: + with open(self._ledger_path(), encoding='utf-8') as fh: + hashes = list(json.load(fh).get('hashes') or []) + except (OSError, ValueError): + pass + self._ledger_order = hashes[-_LEDGER_MAX:] + self._ledger = set(self._ledger_order) + return self._ledger + + def _ledger_delta( + self, msg_dicts: List[Dict[str, + Any]]) -> List[Dict[str, Any]]: + """Conversational messages not yet ingested, in order. Repeated + identical texts dedup to their first occurrence — a repeat carries no + new fact, and it keeps the ledger content-addressed (stable across + context compression rewriting the list).""" + seen = self._ledger_load() + delta: List[Dict[str, Any]] = [] + batch: Set[str] = set() + for m in msg_dicts: + if m.get('role') not in _INGEST_ROLES or not m.get('content'): + continue + h = _content_hash(m) + if h in seen or h in batch: + continue + batch.add(h) + delta.append(m) + return delta + + def _ledger_mark(self, msg_dicts: List[Dict[str, Any]], + persist: bool = True) -> None: + seen = self._ledger_load() + added = False + for m in msg_dicts: + h = _content_hash(m) + if h not in seen: + seen.add(h) + self._ledger_order.append(h) + added = True + if not added or not persist: + return + if len(self._ledger_order) > _LEDGER_MAX: + for stale in self._ledger_order[:-_LEDGER_MAX]: + seen.discard(stale) + self._ledger_order = self._ledger_order[-_LEDGER_MAX:] + try: + os.makedirs(str(self.mem_config.base_dir), exist_ok=True) + fd, tmp = tempfile.mkstemp( + dir=str(self.mem_config.base_dir), suffix='.tmp') + with os.fdopen(fd, 'w', encoding='utf-8') as fh: + json.dump({'version': 1, 'hashes': self._ledger_order}, fh) + os.replace(tmp, self._ledger_path()) + except OSError as e: # pragma: no cover - bookkeeping never breaks + logger.warning(f'[orchestrator] ingest ledger write failed: {e}') # ------------------------------------------------------------------ # Flush (pre-compression) @@ -94,9 +330,10 @@ async def add(self, messages: List[Message], **kwargs: Any) -> None: async def flush(self, messages: List[Message]) -> None: if not self.mem_config.enabled: return - backend = await self._ensure_started() - msg_dicts = _messages_to_dicts(messages) - await backend.on_pre_compress(msg_dicts) + async with _store_lock(self.mem_config.base_dir): + backend = await self._ensure_started() + msg_dicts = _messages_to_dicts(messages) + await backend.on_pre_compress(msg_dicts) # ------------------------------------------------------------------ # Search @@ -148,8 +385,12 @@ def init_update_queue(self) -> None: # ------------------------------------------------------------------ async def close(self) -> None: + # Drain scheduled writes first — closing under a pending ingest would + # either lose the write or race the backend teardown. + await self.flush_pending() if self._backend is not None and self._started: - await self._backend.close() + async with _store_lock(self.mem_config.base_dir): + await self._backend.close() self._started = False # ------------------------------------------------------------------ diff --git a/tests/memory/test_backend_contracts.py b/tests/memory/test_backend_contracts.py index 919aa21c5..7cd30d697 100644 --- a/tests/memory/test_backend_contracts.py +++ b/tests/memory/test_backend_contracts.py @@ -220,6 +220,19 @@ def test_tool_schemas_valid(self, backend_setup): def test_on_messages_no_crash(self, backend_setup): name, backend, loop, tmp = backend_setup + if name == "mem0": + # mem0's ingestion calls a live extraction provider, and the + # adapter now PROPAGATES provider failures by design: the + # swallow-and-report layer moved up to MemoryOrchestrator._ingest, + # which needs the exception to keep failed messages in its delta + # ledger for a retry. A clean provider error is therefore a valid + # outcome here (see tests/memory/test_orchestrator_scheduling.py + # for the orchestrator-level never-raises guarantee). + try: + loop.run_until_complete(backend.on_messages(SAMPLE_TURN)) + except Exception: + pass + return loop.run_until_complete(backend.on_messages(SAMPLE_TURN)) def test_on_pre_compress_no_crash(self, backend_setup): diff --git a/tests/memory/test_orchestrator_scheduling.py b/tests/memory/test_orchestrator_scheduling.py new file mode 100644 index 000000000..01e70f929 --- /dev/null +++ b/tests/memory/test_orchestrator_scheduling.py @@ -0,0 +1,224 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""MemoryOrchestrator write discipline: background scheduling, per-store +serialization, the delta ledger, and status reporting. + +These guard the properties that keep ingestion off the turn's critical path +without corrupting the store or silently losing writes: + +* ingestion is delta-only (content hashes), and a hash is recorded ONLY + after the backend accepted the write — a failure retries by construction; +* retrieval and ingestion serialize on one per-store lock (embedded vector + stores underneath have no locking of their own); +* ``flush_pending`` is a real barrier, so teardown cannot drop the last write; +* ``mark_ingested`` advances the ledger without ingesting (interrupted turns). +""" +import asyncio +import json +import os + +import pytest + +from ms_agent.llm.utils import Message +from ms_agent.memory.unified.config import MemoryConfig +from ms_agent.memory.unified.orchestrator import MemoryOrchestrator + + +class RecordingBackend: + + def __init__(self): + self.batches = [] + self.closed = False + self.active = 0 + self.max_active = 0 + + async def start(self, **kwargs): + pass + + async def on_messages(self, messages, **kwargs): + self.active += 1 + self.max_active = max(self.max_active, self.active) + await asyncio.sleep(0.01) + self.active -= 1 + self.batches.append(list(messages)) + return len(messages) + + async def inject(self, messages): + self.active += 1 + self.max_active = max(self.max_active, self.active) + await asyncio.sleep(0.01) + self.active -= 1 + return messages + + async def on_pre_compress(self, messages): + pass + + async def close(self): + self.closed = True + + def invalidate(self): + pass + + +def _orch(tmp_path, backend, **cfg_kwargs): + orch = MemoryOrchestrator( + MemoryConfig( + base_dir=str(tmp_path), storage_backend='file', **cfg_kwargs)) + orch._backend = backend + orch._started = True + return orch + + +def _msgs(*contents, roles=None): + roles = roles or ['user', 'assistant'] * len(contents) + return [ + Message(role=r, content=c) for r, c in zip(roles, contents) + ] + + +def test_schedule_add_ingests_in_background_and_flush_waits(tmp_path): + + async def main(): + backend = RecordingBackend() + orch = _orch(tmp_path, backend) + task = orch.schedule_add(_msgs('hello', 'hi')) + assert task is not None and not backend.batches # not run inline + await orch.flush_pending() + assert len(backend.batches) == 1 + assert orch.ingest_status['state'] == 'ok' + assert orch.ingest_status['count'] == 2 + + asyncio.run(main()) + + +def test_system_and_tool_rows_never_reach_the_backend(tmp_path): + + async def main(): + backend = RecordingBackend() + orch = _orch(tmp_path, backend) + await orch.add([ + Message(role='system', content='prompt'), + Message(role='user', content='q'), + Message(role='tool', content='{"ok":true}', tool_call_id='1'), + Message(role='assistant', content='a'), + ]) + assert [m['role'] for m in backend.batches[0]] == ['user', 'assistant'] + + asyncio.run(main()) + + +def test_second_ingest_sends_only_the_delta(tmp_path): + + async def main(): + backend = RecordingBackend() + orch = _orch(tmp_path, backend) + history = _msgs('turn one', 'answer one') + await orch.add(history) + history += _msgs('turn two', 'answer two') + await orch.add(history) + assert [m['content'] for m in backend.batches[1]] == [ + 'turn two', 'answer two' + ] + + asyncio.run(main()) + + +def test_ledger_survives_a_process_restart(tmp_path): + + async def main(): + history = _msgs('turn one', 'answer one') + await _orch(tmp_path, RecordingBackend()).add(history) + assert os.path.exists(tmp_path / 'ingest_state.json') + with open(tmp_path / 'ingest_state.json') as fh: + assert len(json.load(fh)['hashes']) == 2 + + # A fresh orchestrator (new process) must not re-ingest old turns. + backend = RecordingBackend() + await _orch(tmp_path, backend).add(history) + assert backend.batches == [] + + asyncio.run(main()) + + +def test_failed_ingest_is_retried_on_the_next_turn(tmp_path): + + async def main(): + + class Failing(RecordingBackend): + + async def on_messages(self, messages, **kwargs): + raise RuntimeError('provider down') + + orch = _orch(tmp_path, Failing()) + await orch.add(_msgs('important fact', 'noted')) + assert orch.ingest_status['state'] == 'error' + assert 'provider down' in orch.ingest_status['error'] + + # Hashes were NOT recorded, so the same messages come back as delta. + backend = RecordingBackend() + orch._backend = backend + await orch.add(_msgs('important fact', 'noted')) + assert [m['content'] for m in backend.batches[0]] == [ + 'important fact', 'noted' + ] + + asyncio.run(main()) + + +def test_mark_ingested_skips_interrupted_content(tmp_path): + + async def main(): + backend = RecordingBackend() + orch = _orch(tmp_path, backend) + interrupted = _msgs('do something', 'half-finished ans') + orch.mark_ingested(interrupted) + await orch.add(interrupted + _msgs('next turn', 'done')) + assert [m['content'] for m in backend.batches[0]] == [ + 'next turn', 'done' + ] + + asyncio.run(main()) + + +def test_ingest_and_inject_serialize_on_the_store_lock(tmp_path): + + async def main(): + backend = RecordingBackend() + orch = _orch(tmp_path, backend) + msgs = _msgs('q', 'a') + orch.schedule_add(msgs) + await orch.run(msgs) # retrieval while the ingest task is pending + await orch.flush_pending() + assert backend.max_active == 1 # never overlapped + + asyncio.run(main()) + + +def test_ingest_interval_batches_turns(tmp_path): + + async def main(): + backend = RecordingBackend() + orch = _orch(tmp_path, backend, ingest_interval=2) + history = _msgs('turn one', 'answer one') + await orch.add(history) + assert backend.batches == [] # skipped: 1 of 2 + history += _msgs('turn two', 'answer two') + await orch.add(history) + # The firing ingest carries everything not yet ingested. + assert [m['content'] for m in backend.batches[0]] == [ + 'turn one', 'answer one', 'turn two', 'answer two' + ] + + asyncio.run(main()) + + +def test_close_drains_pending_then_closes_backend(tmp_path): + + async def main(): + backend = RecordingBackend() + orch = _orch(tmp_path, backend) + orch.schedule_add(_msgs('q', 'a')) + await orch.close() + assert len(backend.batches) == 1 # drained, not dropped + assert backend.closed + + asyncio.run(main()) diff --git a/tests/memory/test_unified_memory.py b/tests/memory/test_unified_memory.py index 6f92cb347..6ce4851bd 100644 --- a/tests/memory/test_unified_memory.py +++ b/tests/memory/test_unified_memory.py @@ -1601,11 +1601,13 @@ def test_on_messages_without_mem0(self): loop.close() def test_invalidate(self): - self.backend._snapshot = "cached" - self.backend._snapshot_dirty = False + # invalidate() must drop the per-turn retrieval cache so the next + # inject re-queries (an external edit changed what search should see). + self.backend._turn_cache_key = "user\x1fquery" + self.backend._turn_cache_results = [{"memory": "cached"}] self.backend.invalidate() - assert self.backend._snapshot is None - assert self.backend._snapshot_dirty is True + assert self.backend._turn_cache_key is None + assert self.backend._turn_cache_results is None def test_close_safe(self): loop = asyncio.new_event_loop() From 41cfac2e0342e3f14c768a8b54760c0e3742088b Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Tue, 11 Aug 2026 02:11:52 +0800 Subject: [PATCH 13/14] Ingest memory on closing rounds only, and give shared stores an owner-side close - add_memory(add_after_step) now fires only when a round closes the turn (assistant reply with no tool calls) and dispatches through the backend's schedule_add when available: tool rounds are intermediate state, and ingesting every round cost O(rounds x history) extraction calls where the closing ingest covers the whole turn. - an interrupted round advances the ingest ledger WITHOUT ingesting (mark_ingested): a half-finished answer is not durable conversational truth and must not be swept into the next turn's delta. - cleanup_tools drains scheduled ingestion (flush only -- memory instances are shared across agents of one store, so closing here would yank the store from a sibling agent); the new SharedMemoryManager.close_matching(base_dir) is the owner-of-last- resort that actually closes instances and releases the embedded store's exclusive file lock. --- ms_agent/agent/llm_agent.py | 62 ++++++++++++++++++++++++++----- ms_agent/memory/memory_manager.py | 30 +++++++++++++++ 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index 1efe70886..ca646b8cf 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -938,6 +938,19 @@ async def cleanup_tools(self): await self.mcp_runtime.stop() if self.tool_manager is not None: await self.tool_manager.cleanup() + # Drain scheduled memory ingestion so a teardown right after the last + # turn cannot lose its write. Flush only — memory instances are shared + # across agents of the same store (SharedMemoryManager), so CLOSING + # them here would yank the store out from under a sibling agent; the + # owner of the shared instance decides when to close (e.g. via + # SharedMemoryManager.close_matching). + for tool in self.memory_tools: + flush = getattr(tool, 'flush_pending', None) + if flush is not None: + try: + await flush(timeout=15) + except Exception as e: # noqa: BLE001 - cleanup is best-effort + logger.warning(f'memory flush on cleanup failed: {e}') @property def stream(self): @@ -1937,6 +1950,20 @@ async def add_memory(self, messages: List[Message], add_type, **kwargs): if not any(v is not None for v in [user_id, agent_id, run_id, memory_type]): continue + # Ingestion runs an extraction LLM + embeddings (seconds). + # Backends that support it take the write off the turn's + # critical path — schedule_add returns immediately and the + # write serializes against retrieval on the store's own lock. + # Others (legacy memory types) keep the awaited behaviour. + if add_type == 'add_after_step' and hasattr( + tool, 'schedule_add'): + tool.schedule_add( + messages, + user_id=user_id, + agent_id=agent_id, + run_id=run_id, + memory_type=memory_type) + continue await tool.add( messages, user_id=user_id, @@ -2250,6 +2277,17 @@ async def run_loop(self, messages: Union[List[Message], str], # plus synthesized interrupted tool results — before the # cancellation unwinds. Sync file I/O only; must re-raise. self._persist_partial_round(messages, pre_step_len) + # An interrupted round is never ingested into long-term + # memory: a half-finished answer is not durable + # conversational truth. Advance the ingest ledger past it + # (sync, in-memory + small file write) so the next turn's + # delta does not sweep the partial content in either. + for _mem_tool in self.memory_tools: + if hasattr(_mem_tool, 'mark_ingested'): + try: + _mem_tool.mark_ingested(messages) + except Exception: # noqa: E722 - never mask cancel + pass raise # Persist THIS round's step output (assistant + any tool @@ -2264,17 +2302,21 @@ async def run_loop(self, messages: Union[List[Message], str], for msg in messages[pre_step_len:step_end_len]: self.session_log.append(self._msg_to_dict(msg)) - # Ingest THIS round's memory here, for the same reason the - # SessionLog write above moved up: after_tool_call blocks on the - # next prompt in interactive mode. Running afterwards made - # 'add_after_step' mean "after the *next* step" — round N was + # Ingest memory here — before after_tool_call, which blocks on + # the next prompt in interactive mode (running afterwards made + # 'add_after_step' mean "after the *next* step": round N was # only ingested once round N+1 arrived, a single-round session - # was never ingested at all (the run-loop tail that would have - # caught it is skipped when the input source raises EOF), and - # the message list it saw had the next user turn already - # appended. Before the block, `messages` is exactly this round. - await self.add_memory( - messages, add_type='add_after_step', **kwargs) + # never at all, and the ingested list already carried the next + # user turn). Only on a CLOSING round — an assistant reply with + # no tool calls, i.e. the turn is complete. Tool rounds are + # intermediate state, not durable conversational truth, and + # ingesting every round made memory cost O(rounds x history): + # a 4-round tool-calling answer paid ~4 extraction-LLM calls + # where one covers it (the closing ingest sees the whole turn). + if (messages and messages[-1].role == 'assistant' + and not messages[-1].tool_calls): + await self.add_memory( + messages, add_type='add_after_step', **kwargs) await self.after_tool_call(messages) self.runtime.round += 1 diff --git a/ms_agent/memory/memory_manager.py b/ms_agent/memory/memory_manager.py index 14d7afb3e..77b1ce312 100644 --- a/ms_agent/memory/memory_manager.py +++ b/ms_agent/memory/memory_manager.py @@ -44,6 +44,36 @@ async def get_shared_memory(cls, config: DictConfig, return cls._instances[key] + @classmethod + async def close_matching(cls, base_dir: str) -> int: + """Close and drop every shared instance rooted at ``base_dir``. + + The owner-of-last-resort for embedded stores: closing an instance + releases its vector client (and with it the store's exclusive file + lock), which per-agent cleanup deliberately does NOT do — an + instance may be shared by several live agents, so only whoever knows + no agent still needs the store (e.g. a runtime registry evicting the + last session of a project) may call this. Returns how many instances + were closed.""" + target = os.path.abspath(os.path.expanduser(str(base_dir))) + closed = 0 + for key, mem in list(cls._instances.items()): + mem_cfg = getattr(mem, 'mem_config', None) + base = getattr(mem_cfg, 'base_dir', None) + if base is None or os.path.abspath(str(base)) != target: + continue + try: + close = getattr(mem, 'close', None) + if close is not None: + await close() + except Exception as e: # noqa: BLE001 - eviction is best-effort + logger.warning( + f'closing shared memory for {key} failed: {e}') + cls._instances.pop(key, None) + closed += 1 + logger.info(f'Closed shared memory instance: {key}') + return closed + @classmethod def clear_shared_memory(cls, config: DictConfig, mem_instance_type: str): """Clear shared memory instances. If config is provided, clear specific instance.""" From 313a2327596a51e9f45fda9940428945954f06af Mon Sep 17 00:00:00 2001 From: alcholiclg Date: Tue, 11 Aug 2026 13:49:49 +0800 Subject: [PATCH 14/14] Make mem0 recall size configurable (recall_top_k) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The number of recalled memories injected per turn was hardcoded twice (search default 20, then a [:10] formatting slice). MemoryConfig gains recall_top_k (default 10, read from the unified_memory node) and the mem0 adapter threads it through search and formatting — consumers can now size recall to their context budget. --- ms_agent/memory/unified/backends/mem0_adapter.py | 16 +++++++++------- ms_agent/memory/unified/config.py | 7 ++++++- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/ms_agent/memory/unified/backends/mem0_adapter.py b/ms_agent/memory/unified/backends/mem0_adapter.py index 4323763aa..0d8c94410 100644 --- a/ms_agent/memory/unified/backends/mem0_adapter.py +++ b/ms_agent/memory/unified/backends/mem0_adapter.py @@ -46,12 +46,12 @@ def _result_list(results: Any) -> List[Dict[str, Any]]: return list(results or []) -def _mem0_search(m0: Any, query: str, user_id: str) -> Any: +def _mem0_search(m0: Any, query: str, user_id: str, top_k: int = 10) -> Any: """mem0 2.x moved entity params into ``filters=``; 1.x uses kwargs.""" try: - return m0.search(query, filters={'user_id': user_id}) + return m0.search(query, filters={'user_id': user_id}, top_k=top_k) except TypeError: - return m0.search(query, user_id=user_id) + return m0.search(query, user_id=user_id, limit=top_k) class Mem0Backend(BaseMemoryBackend): @@ -121,10 +121,11 @@ async def inject( and self._turn_cache_results is not None: results = self._turn_cache_results else: + top_k = max(1, int(getattr(self._config, 'recall_top_k', 10))) try: results = _result_list( await _offload(_mem0_search, self._mem0, query, - self._user_id)) + self._user_id, top_k)) except Exception as e: logger.debug(f'[mem0_backend] search failed: {e}') return messages @@ -133,7 +134,8 @@ async def inject( if not results: return messages - formatted = self._format_results(results) + formatted = self._format_results( + results, max(1, int(getattr(self._config, 'recall_top_k', 10)))) if not formatted: return messages @@ -220,9 +222,9 @@ def _extract_query(messages: List[Dict[str, Any]]) -> str: return '' @staticmethod - def _format_results(results: Any) -> str: + def _format_results(results: Any, top_k: int = 10) -> str: lines = [] - for r in _result_list(results)[:10]: + for r in _result_list(results)[:top_k]: text = r.get('memory', r.get('text', '')) if text: lines.append(f'- {text}') diff --git a/ms_agent/memory/unified/config.py b/ms_agent/memory/unified/config.py index 5b57cde2e..0640bb64b 100644 --- a/ms_agent/memory/unified/config.py +++ b/ms_agent/memory/unified/config.py @@ -31,6 +31,10 @@ class MemoryConfig: # everything not yet ingested on the next firing ingest. ingest_interval: int = 1 + # How many recalled memories retrieval-style backends (mem0) inject per + # turn. + recall_top_k: int = 10 + # Backend-specific options keyed by backend name backend_options: Dict[str, Any] = field(default_factory=dict) @@ -118,7 +122,8 @@ def from_dict_config(cls, cfg: DictConfig) -> 'MemoryConfig': if k in ns: flat[k] = ns[k] - for k in ('enabled', 'base_dir', 'llm_config', 'ingest_interval'): + for k in ('enabled', 'base_dir', 'llm_config', 'ingest_interval', + 'recall_top_k'): if k in raw: flat[k] = raw[k]