diff --git a/agent/memory_manager.py b/agent/memory_manager.py index 474d9b3ac9..b50b22a103 100644 --- a/agent/memory_manager.py +++ b/agent/memory_manager.py @@ -46,28 +46,21 @@ def _has_var_kwargs(params) -> bool: def _accepts_require_checkpoint(fn: Callable[..., Any]) -> bool: - """True if ``fn`` can receive the ``require_checkpoint`` keyword. + """True if ``fn`` can receive the ``require_checkpoint`` keyword (unreadable signatures -> False). - v2 providers written against the docs example use the bare ``on_pre_compress(self, - messages)`` shape; passing the keyword would raise TypeError, which the host would - re-raise as a checkpoint failure even though the durable write succeeded. - Unreadable signatures conservatively report False. + Bare-shape v2 providers (``on_pre_compress(self, messages)``) would raise TypeError on the + keyword, which the host would re-raise as a checkpoint failure despite a successful write. """ params = _signature_params(fn) if params is None: return False - param = params.get("require_checkpoint") - return _has_var_kwargs(params) or ( - param is not None and param.kind in (inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) - ) + kind = getattr(params.get("require_checkpoint"), "kind", None) + return _has_var_kwargs(params) or kind in (inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) def _ctx_bound(fn: Callable[[], Any]) -> Callable[[], Any]: - """Bind ``fn`` to the CALLER's contextvars for execution on another thread. - - Profile isolation is a ContextVar-scoped HERMES_HOME override; an unbound provider - on a worker thread would silently resolve paths/secrets against the default profile. - """ + """Bind ``fn`` to the CALLER's contextvars for another thread: profile isolation is a + ContextVar-scoped HERMES_HOME override, and an unbound worker would silently use the default profile.""" return partial(contextvars.copy_context().run, fn) @@ -88,10 +81,8 @@ def normalize_tool_schema(schema: Any) -> Optional[Dict[str, Any]]: return schema if name and isinstance(name, str) else None -def memory_provider_tools_enabled( - enabled_toolsets: Optional[List[str]], disabled_toolsets: Optional[List[str]] = None, *, - memory_tool_present: bool = False, -) -> bool: +def memory_provider_tools_enabled(enabled_toolsets: Optional[List[str]], disabled_toolsets: Optional[List[str]] = None, + *, memory_tool_present: bool = False) -> bool: """Return whether external memory-provider tools should be exposed.""" if disabled_toolsets and "memory" in disabled_toolsets: return False @@ -121,11 +112,9 @@ def memory_provider_tools_exposed(agent: Any) -> bool: never advertises tools absent from the tool surface. """ tools = getattr(agent, "tools", None) - return memory_provider_tools_enabled( - getattr(agent, "enabled_toolsets", None), - getattr(agent, "disabled_toolsets", None), - memory_tool_present=isinstance(tools, (list, tuple)) and any(_tool_name(t) == "memory" for t in tools), - ) + present = isinstance(tools, (list, tuple)) and any(_tool_name(t) == "memory" for t in tools) + enabled, disabled = getattr(agent, "enabled_toolsets", None), getattr(agent, "disabled_toolsets", None) + return memory_provider_tools_enabled(enabled, disabled, memory_tool_present=present) def inject_memory_provider_tools(agent: Any) -> int: @@ -138,7 +127,8 @@ def inject_memory_provider_tools(agent: Any) -> int: if not memory_provider_tools_exposed(agent): # Say so once: a silent 0 leaves the provider looking "half on" with no clue which # config key (platform_toolsets / disabled_toolsets) gated it. - _providers = [p for p in (getattr(memory_manager, "providers", None) or []) if getattr(p, "name", "") != "builtin"] + _providers = [p for p in getattr(memory_manager, "providers", None) or [] + if getattr(p, "name", "") != "builtin"] if _providers: logger.info( "Memory provider(s) %s configured but the 'memory' toolset is " @@ -192,10 +182,9 @@ def sanitize_context(text: str) -> str: class StreamingContextScrubber: """Stateful scrubber for streaming text whose memory-context spans may straddle deltas. - ``sanitize_context`` needs both tags in one string, so a span split across deltas - would leak to the UI. This holds back partial-tag tails between ``feed()`` calls and - drops everything inside a span. One scrubber (or ``reset()``) per top-level - response; call ``flush()`` at end of stream. + ``sanitize_context`` needs both tags in one string, so a split span would leak to the UI; + this holds back partial-tag tails between ``feed()`` calls and drops span interiors. + One scrubber (or ``reset()``) per top-level response; call ``flush()`` at end of stream. """ _OPEN_TAG = "" @@ -218,13 +207,16 @@ class StreamingContextScrubber: out: list[str] = [] while buf: if self._in_span: - idx = buf.lower().find(self._CLOSE_TAG) - held = self._max_partial_suffix(buf, self._CLOSE_TAG) # potential partial close tag tag = self._CLOSE_TAG + idx = buf.lower().find(tag) + held = self._max_partial_suffix(buf, tag) # potential partial close tag else: - idx = self._find_boundary_open_tag(buf) - held = self._max_pending_open_suffix(buf) or self._max_partial_suffix(buf, self._OPEN_TAG) tag = self._OPEN_TAG + idx = self._find_boundary_open_tag(buf) + # A complete boundary tag at the buffer end is held until the next char confirms it. + n = len(tag) + pending = n if buf.lower().endswith(tag) and self._ends_at_block_boundary(buf[:-n]) else 0 + held = pending or self._max_partial_suffix(buf, tag) if idx == -1: # Hold back the possible partial tag; inside a span the rest is dropped. if not self._in_span: @@ -248,41 +240,26 @@ class StreamingContextScrubber: @staticmethod def _max_partial_suffix(buf: str, tag: str) -> int: """Length of the longest buf-suffix that is a (case-insensitive) prefix of ``tag``, else 0.""" - tag_lower = tag.lower() - buf_lower = buf.lower() - for i in range(min(len(buf_lower), len(tag_lower) - 1), 0, -1): - if tag_lower.startswith(buf_lower[-i:]): - return i - return 0 + tag_lower, buf_lower = tag.lower(), buf.lower() + span = range(min(len(buf_lower), len(tag_lower) - 1), 0, -1) + return next((i for i in span if tag_lower.startswith(buf_lower[-i:])), 0) def _find_boundary_open_tag(self, buf: str) -> int: """Find an opening fence only when it starts a block-like span (own line, newline after).""" - buf_lower = buf.lower() - search_start = 0 - while True: - idx = buf_lower.find(self._OPEN_TAG, search_start) - if idx == -1: - return -1 - after_idx = idx + len(self._OPEN_TAG) - if self._is_block_boundary(buf, idx) and after_idx < len(buf) and buf[after_idx] in "\r\n": + buf_lower, tag_len = buf.lower(), len(self._OPEN_TAG) + idx = buf_lower.find(self._OPEN_TAG) + while idx != -1: + after_idx = idx + tag_len + if self._ends_at_block_boundary(buf[:idx]) and after_idx < len(buf) and buf[after_idx] in "\r\n": return idx - search_start = idx + 1 - - def _max_pending_open_suffix(self, buf: str) -> int: - """Hold a complete boundary tag at the buffer end until the following char confirms it.""" - if buf.lower().endswith(self._OPEN_TAG) and self._is_block_boundary(buf, len(buf) - len(self._OPEN_TAG)): - return len(self._OPEN_TAG) - return 0 + idx = buf_lower.find(self._OPEN_TAG, idx + 1) + return -1 def _ends_at_block_boundary(self, text: str) -> bool: - """Whether emitting ``text`` leaves the stream at a line start (blank tail after the last newline).""" - last_newline = text.rfind("\n") - if last_newline == -1: - return self._at_block_boundary and text.strip() == "" - return text[last_newline + 1:].strip() == "" - - def _is_block_boundary(self, buf: str, idx: int) -> bool: - return self._at_block_boundary if idx == 0 else self._ends_at_block_boundary(buf[:idx]) + """Whether emitting ``text`` leaves the stream at a line start (blank tail after the last newline; + no newline at all -> only whitespace and already at a boundary).""" + head, sep, tail = text.rpartition("\n") + return tail.strip() == "" and (bool(sep) or self._at_block_boundary) def _append_visible(self, out: list[str], text: str) -> None: if text: @@ -307,11 +284,6 @@ def build_memory_context_block(raw_context: str) -> str: ) -def _nonblank(text: Any) -> Any: - """Return ``text`` when it has non-whitespace content, else None.""" - return text if text and text.strip() else None - - class MemoryManager: """Builtin provider (always first) plus at most one external provider. @@ -323,11 +295,11 @@ class MemoryManager: self._providers: List[MemoryProvider] = [] self._tool_to_provider: Dict[str, MemoryProvider] = {} self._has_external: bool = False - self._external_prefetch_timeout = ( - _EXTERNAL_PREFETCH_TIMEOUT_S if external_prefetch_timeout is None else float(external_prefetch_timeout) - ) - if self._external_prefetch_timeout <= 0: + timeout = external_prefetch_timeout + timeout = _EXTERNAL_PREFETCH_TIMEOUT_S if timeout is None else float(timeout) + if timeout <= 0: raise ValueError("external_prefetch_timeout must be positive") + self._external_prefetch_timeout = timeout self._external_prefetch_threads: Dict[str, threading.Thread] = {} self._external_prefetch_lock = threading.Lock() # Single-worker background executor for end-of-turn sync/prefetch, created lazily so @@ -342,15 +314,10 @@ class MemoryManager: "status": "not_started", "abandoned_writes": 0, "abandoned_prefetches": 0, "active_tasks": 0, } - def _each_provider( - self, label: str, call: Callable[[MemoryProvider], Any], *, - level: int = logging.DEBUG, providers: Optional[List[MemoryProvider]] = None, exc_info: bool = False, - ) -> List[Any]: - """Call ``call(provider)`` for each provider, logging and swallowing failures. - - ``label`` completes the log line ``Memory provider ''