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 423126f19..ca646b8cf 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. @@ -901,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): @@ -1727,37 +1777,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 @@ -1926,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, @@ -2239,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 @@ -2253,6 +2302,22 @@ 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 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 + # 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 @@ -2263,9 +2328,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 diff --git a/ms_agent/llm/router.py b/ms_agent/llm/router.py index b8dc97102..54a38954f 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)`` @@ -107,7 +107,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/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.""" diff --git a/ms_agent/memory/unified/backends/mem0_adapter.py b/ms_agent/memory/unified/backends/mem0_adapter.py index e56740b8f..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): @@ -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 ──────────────────────────────────────────────────── @@ -84,6 +89,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 ─────────────────────────────────────────────────────── @@ -99,16 +116,26 @@ 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: + 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, top_k)) + 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) + formatted = self._format_results( + results, max(1, int(getattr(self._config, 'recall_top_k', 10)))) if not formatted: return messages @@ -127,24 +154,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: - logger.warning(f'[mem0_backend] add failed: {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 ─────────────────────────────────────────────────────── @@ -172,8 +207,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 ───────────────────────────────────────────── @@ -186,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 7fa247e5b..0640bb64b 100644 --- a/ms_agent/memory/unified/config.py +++ b/ms_agent/memory/unified/config.py @@ -26,6 +26,15 @@ 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 + + # 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) @@ -113,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'): + for k in ('enabled', 'base_dir', 'llm_config', 'ingest_interval', + 'recall_top_k'): 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/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/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/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') 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() 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 = [] 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)