unique-user-memory 2026.32.0.dev1__tar.gz → 2026.32.0.dev3__tar.gz

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev3}/PKG-INFO

@@ -1,14 +1,14 @@
 Metadata-Version: 2.3
 Name: unique-user-memory
-Version: 2026.32.0.dev1
+Version: 2026.32.0.dev3
 Summary: 
 Author: Fabian Schläpfer
 Author-email: Fabian Schläpfer <fabian@unique.ch>
 License: Proprietary
 Requires-Dist: jinja2>=3.1.6
 Requires-Dist: pydantic>=2.8.2
-Requires-Dist: unique-sdk>=2026.32.0.dev0,<2026.32.0rc0
-Requires-Dist: unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-sdk>=2026.32.0.dev8,<2026.32.0rc0
+Requires-Dist: unique-toolkit>=2026.32.0.dev8,<2026.32.0rc0
 Requires-Python: >=3.12, <4
 Description-Content-Type: text/markdown
 
@@ -23,7 +23,8 @@ Persistent per-user memory for Unique AI agents.
 The package provides:
 
 - `UserMemoryConfig` - Pydantic configuration for the consolidation model, profile token budget, and memory folder.
-- `load_user_memory(...)` - resolves the user's private memory folder, downloads `memory.md`, and enforces the configured token budget. The `language_model` argument is used to tokenize `memory.md` when capping it, so it must be the same effective model the postprocessor uses for consolidation (see Integration below).
+- `load_user_memory(...)` - resolves the user's private memory folder, downloads `memory.md`, and enforces the configured token budget. The `language_model` argument is used to tokenize `memory.md` when capping it, so it must be the same effective model the postprocessor uses for consolidation (see Integration below). Returns a `UserMemoryState` with the profile text and scope id.
+- `UserMemoryMessageLogger` - emits chat Steps (MessageLogs) for load and update, including typed `UserMemory` detail entries the chat frontend renders as a badge that opens Settings → Context Memory. Frontends that do not know the entry type render nothing, so the entries are safe to emit in any deploy order.
 - `UserMemoryPostprocessor` - runs after the assistant response, consolidates the latest turn into the profile, and uploads the updated `memory.md`.
 
 The memory file is intentionally small and structured. It is rewritten as a full Markdown profile rather than appended to as an event log.
@@ -31,11 +32,13 @@ The memory file is intentionally small and structured. It is rewritten as a full
 ## Lifecycle
 
 1. The orchestrator enables memory when `space.allow_user_memory` is true.
-2. `load_user_memory(...)` resolves the pre-provisioned root folder, ensures a private child folder for the current user, and downloads `/user-memory/<user_id>/memory.md` if it exists.
-3. The loaded memory text is passed into the agent context for the current turn.
-4. `UserMemoryPostprocessor` runs after the assistant response.
-5. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
-6. If the profile changed, `memory.md` is uploaded back to the user's folder with ingestion skipped and the content hidden from chat.
+2. The orchestrator emits a **Loading context memory** Step, then `load_user_memory(...)` resolves the pre-provisioned root folder, ensures a private child folder for the current user, and downloads `/user-memory/<user_id>/memory.md` if it exists.
+3. When load returns a `UserMemoryState`, that Step is completed with a **Context memory** detail entry (`type: UserMemory`) that the chat frontend renders as a badge opening Settings → Context Memory. A successful `None` return (soft skip) completes the Step without the entry; a raised exception marks the Step failed.
+4. If memory was loaded, its text is passed into the agent context for the current turn.
+5. `UserMemoryPostprocessor` runs after the assistant response.
+6. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
+7. If a rewrite runs, an **Updating your memory** Step is shown while consolidating (no settings entry yet).
+8. If the profile changed and `memory.md` uploads successfully (ingestion skipped, content hidden from chat), that Step is completed with a **Review your context memory** detail entry (same settings badge). On NOOP or failed upload the Step completes without the entry.
 
 ## Storage Model
 
@@ -109,7 +112,9 @@ Typical orchestration code loads memory before the agent loop and registers the
 `load_user_memory` and `UserMemoryPostprocessor` must be given the **same** effective language model: the postprocessor consolidates memory with either the orchestrator model or the configured one depending on `use_orchestrator_language_model`, and load-time token capping must use that same model so the loaded baseline is tokenized the way consolidation expects. Resolve the effective model once and pass it to both:
 
 ```python
+from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
 from unique_user_memory.user_memory import load_user_memory
+from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
 from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
 
 user_memory_config = config.agent.services.user_memory_config
@@ -122,14 +127,36 @@ memory_language_model = (
     else user_memory_config.language_model
 )
 
-user_memory_state = await load_user_memory(
-    event=event,
-    config=user_memory_config,
-    language_model=memory_language_model,
+message_step_logger = MessageStepLogger(chat_service)
+memory_message_step_logger = UserMemoryMessageLogger(
+    message_step_logger,
     logger=logger,
 )
-
-if user_memory_state is not None:
+await memory_message_step_logger.log_loading_start()
+user_memory_state = None
+load_succeeded = False
+try:
+    user_memory_state = await load_user_memory(
+        event=event,
+        config=user_memory_config,
+        language_model=memory_language_model,
+        logger=logger,
+    )
+    load_succeeded = True
+except Exception as exc:
+    logger.warning(
+        "[user-memory] load raised - running without memory: [%s] %s",
+        type(exc).__name__,
+        exc,
+    )
+finally:
+    # Always close the RUNNING step — otherwise the chat Steps UI stays stuck
+    # on "Loading context memory" for that turn when load raises.
+    if not load_succeeded:
+        await memory_message_step_logger.log_loading_failed()
+
+if load_succeeded and user_memory_state is not None:
+    await memory_message_step_logger.log_loading_complete(with_settings_entry=True)
     user_memory_text = user_memory_state.text
     postprocessor_manager.add_postprocessor(
         UserMemoryPostprocessor(
@@ -138,8 +165,11 @@ if user_memory_state is not None:
             event=event,
             state=user_memory_state,
             logger=logger,
+            message_step_logger=memory_message_step_logger,
         )
     )
+elif load_succeeded:
+    await memory_message_step_logger.log_loading_complete(with_settings_entry=False)
 ```
 
 Note that `UserMemoryPostprocessor` re-derives the effective model internally from `use_orchestrator_language_model`, so passing `memory_language_model` (rather than the raw orchestrator model) keeps its behavior identical while ensuring `load_user_memory` caps with the matching tokenizer.

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev3}/README.md

@@ -9,7 +9,8 @@ Persistent per-user memory for Unique AI agents.
 The package provides:
 
 - `UserMemoryConfig` - Pydantic configuration for the consolidation model, profile token budget, and memory folder.
-- `load_user_memory(...)` - resolves the user's private memory folder, downloads `memory.md`, and enforces the configured token budget. The `language_model` argument is used to tokenize `memory.md` when capping it, so it must be the same effective model the postprocessor uses for consolidation (see Integration below).
+- `load_user_memory(...)` - resolves the user's private memory folder, downloads `memory.md`, and enforces the configured token budget. The `language_model` argument is used to tokenize `memory.md` when capping it, so it must be the same effective model the postprocessor uses for consolidation (see Integration below). Returns a `UserMemoryState` with the profile text and scope id.
+- `UserMemoryMessageLogger` - emits chat Steps (MessageLogs) for load and update, including typed `UserMemory` detail entries the chat frontend renders as a badge that opens Settings → Context Memory. Frontends that do not know the entry type render nothing, so the entries are safe to emit in any deploy order.
 - `UserMemoryPostprocessor` - runs after the assistant response, consolidates the latest turn into the profile, and uploads the updated `memory.md`.
 
 The memory file is intentionally small and structured. It is rewritten as a full Markdown profile rather than appended to as an event log.
@@ -17,11 +18,13 @@ The memory file is intentionally small and structured. It is rewritten as a full
 ## Lifecycle
 
 1. The orchestrator enables memory when `space.allow_user_memory` is true.
-2. `load_user_memory(...)` resolves the pre-provisioned root folder, ensures a private child folder for the current user, and downloads `/user-memory/<user_id>/memory.md` if it exists.
-3. The loaded memory text is passed into the agent context for the current turn.
-4. `UserMemoryPostprocessor` runs after the assistant response.
-5. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
-6. If the profile changed, `memory.md` is uploaded back to the user's folder with ingestion skipped and the content hidden from chat.
+2. The orchestrator emits a **Loading context memory** Step, then `load_user_memory(...)` resolves the pre-provisioned root folder, ensures a private child folder for the current user, and downloads `/user-memory/<user_id>/memory.md` if it exists.
+3. When load returns a `UserMemoryState`, that Step is completed with a **Context memory** detail entry (`type: UserMemory`) that the chat frontend renders as a badge opening Settings → Context Memory. A successful `None` return (soft skip) completes the Step without the entry; a raised exception marks the Step failed.
+4. If memory was loaded, its text is passed into the agent context for the current turn.
+5. `UserMemoryPostprocessor` runs after the assistant response.
+6. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
+7. If a rewrite runs, an **Updating your memory** Step is shown while consolidating (no settings entry yet).
+8. If the profile changed and `memory.md` uploads successfully (ingestion skipped, content hidden from chat), that Step is completed with a **Review your context memory** detail entry (same settings badge). On NOOP or failed upload the Step completes without the entry.
 
 ## Storage Model
 
@@ -95,7 +98,9 @@ Typical orchestration code loads memory before the agent loop and registers the
 `load_user_memory` and `UserMemoryPostprocessor` must be given the **same** effective language model: the postprocessor consolidates memory with either the orchestrator model or the configured one depending on `use_orchestrator_language_model`, and load-time token capping must use that same model so the loaded baseline is tokenized the way consolidation expects. Resolve the effective model once and pass it to both:
 
 ```python
+from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
 from unique_user_memory.user_memory import load_user_memory
+from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
 from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
 
 user_memory_config = config.agent.services.user_memory_config
@@ -108,14 +113,36 @@ memory_language_model = (
     else user_memory_config.language_model
 )
 
-user_memory_state = await load_user_memory(
-    event=event,
-    config=user_memory_config,
-    language_model=memory_language_model,
+message_step_logger = MessageStepLogger(chat_service)
+memory_message_step_logger = UserMemoryMessageLogger(
+    message_step_logger,
     logger=logger,
 )
-
-if user_memory_state is not None:
+await memory_message_step_logger.log_loading_start()
+user_memory_state = None
+load_succeeded = False
+try:
+    user_memory_state = await load_user_memory(
+        event=event,
+        config=user_memory_config,
+        language_model=memory_language_model,
+        logger=logger,
+    )
+    load_succeeded = True
+except Exception as exc:
+    logger.warning(
+        "[user-memory] load raised - running without memory: [%s] %s",
+        type(exc).__name__,
+        exc,
+    )
+finally:
+    # Always close the RUNNING step — otherwise the chat Steps UI stays stuck
+    # on "Loading context memory" for that turn when load raises.
+    if not load_succeeded:
+        await memory_message_step_logger.log_loading_failed()
+
+if load_succeeded and user_memory_state is not None:
+    await memory_message_step_logger.log_loading_complete(with_settings_entry=True)
     user_memory_text = user_memory_state.text
     postprocessor_manager.add_postprocessor(
         UserMemoryPostprocessor(
@@ -124,8 +151,11 @@ if user_memory_state is not None:
             event=event,
             state=user_memory_state,
             logger=logger,
+            message_step_logger=memory_message_step_logger,
         )
     )
+elif load_succeeded:
+    await memory_message_step_logger.log_loading_complete(with_settings_entry=False)
 ```
 
 Note that `UserMemoryPostprocessor` re-derives the effective model internally from `use_orchestrator_language_model`, so passing `memory_language_model` (rather than the raw orchestrator model) keeps its behavior identical while ensuring `load_user_memory` caps with the matching tokenizer.

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_user_memory"
-version = "2026.32.0.dev1"
+version = "2026.32.0.dev3"
 description = ""
 authors = [
     { name = "Fabian Schläpfer", email = "fabian@unique.ch" },
@@ -11,8 +11,8 @@ requires-python = ">=3.12,<4"
 dependencies = [
     "jinja2>=3.1.6",
     "pydantic>=2.8.2",
-    "unique-sdk>=2026.32.0.dev0,<2026.32.0rc0",
-    "unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-sdk>=2026.32.0.dev8,<2026.32.0rc0",
+    "unique-toolkit>=2026.32.0.dev8,<2026.32.0rc0",
 ]
 
 [dependency-groups]

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev3}/unique_user_memory/config.py

@@ -42,14 +42,6 @@ class UserMemoryConfig(BaseModel):
             "runs."
         ),
     )
-    updating_notice_enabled: bool = Field(
-        default=True,
-        description=(
-            "When true, a transient 'updating context memory' notice is appended to the "
-            "assistant message while the memory rewrite runs, and removed "
-            "again once it completes."
-        ),
-    )
     root_folder: Annotated[str, RJSFMetaTag.SpecialWidget.hidden()] = Field(
         default="user-memory",
         min_length=1,

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev3}/unique_user_memory/tests/test_user_memory.py

@@ -2,10 +2,13 @@ import asyncio
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from unique_toolkit.chat.schemas import MessageLogStatus
 from unique_toolkit.language_model.default_language_model import (
     DEFAULT_LANGUAGE_MODEL,
 )
 from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
+from unique_toolkit.language_model.schemas import LanguageModelTokenUsage
 
 from unique_user_memory.config import UserMemoryConfig
 from unique_user_memory.user_memory import (
@@ -18,10 +21,10 @@ from unique_user_memory.user_memory import (
     enforce_token_cap,
     ensure_user_memory_folder,
     fit_user_memory,
-    noop_update_callback,
     should_consolidate_memory,
     upload_user_memory,
 )
+from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
 from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
 from unique_user_memory.user_memory_prompts import (
     consolidation_system_prompt,
@@ -677,67 +680,113 @@ async def test_consolidate_user_memory_invokes_update_end_when_start_cancelled(
 
 
 @pytest.mark.asyncio
-async def test_user_memory_postprocessor_shows_and_removes_updating_notice(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    updated_memory = "# User Memory\n\n## Identity\n- Updated"
-    original_text = "Here is your answer."
+async def test_user_memory_message_step_logger_load_emits_settings_detail_entry() -> (
+    None
+):
+    step_logger = MagicMock()
+    step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
+    message_logger = UserMemoryMessageLogger(step_logger)
+
+    await message_logger.log_loading_start()
+    await message_logger.log_loading_complete(with_settings_entry=True)
+
+    assert step_logger.create_or_update_message_log_async.await_count == 2
+    start_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
+        0
+    ].kwargs
+    assert start_kwargs["header"] == "Loading context memory"
+    assert start_kwargs["details"] is None
+    assert start_kwargs["references"] == []
+    complete_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
+        1
+    ].kwargs
+    assert complete_kwargs["status"] == MessageLogStatus.COMPLETED
+    assert complete_kwargs["references"] == []
+    entry = complete_kwargs["details"].data[0]
+    assert entry.type == "UserMemory"
+    assert entry.text == "Context memory"
 
-    async def fake_consolidate(*, on_update_start, on_update_end, **kwargs) -> str:  # type: ignore[no-untyped-def]
-        await on_update_start()
-        await on_update_end()
-        return updated_memory
 
-    monkeypatch.setattr(
-        "unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
-        fake_consolidate,
-    )
-    monkeypatch.setattr(
-        "unique_user_memory.user_memory_postprocessor.upload_user_memory",
-        AsyncMock(return_value=True),
-    )
-    chat_service = MagicMock()
-    chat_service.modify_assistant_message_async = AsyncMock()
-    monkeypatch.setattr(
-        "unique_user_memory.user_memory_postprocessor.ChatService",
-        MagicMock(return_value=chat_service),
-    )
-    event = MagicMock()
-    event.user_id = "user_1"
-    event.company_id = "company_1"
-    event.payload.user_message.text = "remember this"
-    loop_response = MagicMock()
-    loop_response.message.text = original_text
-    loop_response.message.id = "msg_1"
-    loop_response.message.references = []
-    postprocessor = UserMemoryPostprocessor(
-        config=UserMemoryConfig(),
-        language_model=_TEST_LANGUAGE_MODEL,
-        event=event,
-        state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
-        logger=MagicMock(),
-        chat_service=chat_service,
-    )
+@pytest.mark.asyncio
+async def test_user_memory_message_step_logger_load_skips_entry_when_disabled() -> None:
+    step_logger = MagicMock()
+    step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
+    message_logger = UserMemoryMessageLogger(step_logger)
 
-    await postprocessor.run(loop_response)
+    await message_logger.log_loading_complete(with_settings_entry=False)
+
+    complete_kwargs = step_logger.create_or_update_message_log_async.await_args.kwargs
+    assert complete_kwargs["details"] is None
+    assert complete_kwargs["references"] == []
 
-    calls = chat_service.modify_assistant_message_async.await_args_list
-    assert len(calls) == 2
-    assert calls[0].kwargs["content"].startswith(original_text)
-    assert calls[0].kwargs["content"] != original_text
-    assert calls[0].kwargs["message_id"] == "msg_1"
-    assert calls[1].kwargs["content"] == original_text
+
+@pytest.mark.asyncio
+async def test_user_memory_message_step_logger_load_failed_marks_failed_status() -> (
+    None
+):
+    """
+    Purpose: log_loading_failed updates the loading Step to FAILED with no
+    settings entry.
+    Why this matters: Callers must close a RUNNING loading Step when load
+    raises so the chat UI does not stay stuck.
+    Setup summary: Start then fail the loading step; assert FAILED status.
+    """
+    step_logger = MagicMock()
+    step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
+    message_logger = UserMemoryMessageLogger(step_logger)
+
+    await message_logger.log_loading_start()
+    await message_logger.log_loading_failed()
+
+    failed_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
+        -1
+    ].kwargs
+    assert failed_kwargs["header"] == "Loading context memory"
+    assert failed_kwargs["status"] == MessageLogStatus.FAILED
+    assert failed_kwargs["details"] is None
+    assert failed_kwargs["references"] == []
 
 
 @pytest.mark.asyncio
-async def test_user_memory_postprocessor_skips_notice_when_disabled(
+async def test_user_memory_message_step_logger_update_attaches_review_entry_only_when_requested() -> (
+    None
+):
+    step_logger = MagicMock()
+    step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
+    message_logger = UserMemoryMessageLogger(step_logger)
+
+    await message_logger.log_updating_start()
+    await message_logger.log_updating_complete(with_settings_entry=False)
+    await message_logger.log_updating_complete(with_settings_entry=True)
+
+    assert step_logger.create_or_update_message_log_async.await_count == 3
+    start_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
+        0
+    ].kwargs
+    assert start_kwargs["header"] == "Updating your memory"
+    assert start_kwargs["details"] is None
+    complete_without_entry = (
+        step_logger.create_or_update_message_log_async.await_args_list[1].kwargs
+    )
+    assert complete_without_entry["details"] is None
+    complete_with_entry = (
+        step_logger.create_or_update_message_log_async.await_args_list[2].kwargs
+    )
+    assert complete_with_entry["references"] == []
+    entry = complete_with_entry["details"].data[0]
+    assert entry.type == "UserMemory"
+    assert entry.text == "Review your context memory"
+
+
+@pytest.mark.asyncio
+async def test_user_memory_postprocessor_emits_updating_message_logs(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     updated_memory = "# User Memory\n\n## Identity\n- Updated"
 
     async def fake_consolidate(*, on_update_start, on_update_end, **kwargs) -> str:  # type: ignore[no-untyped-def]
-        assert on_update_start is noop_update_callback
-        assert on_update_end is noop_update_callback
+        await on_update_start()
+        await on_update_end()
         return updated_memory
 
     monkeypatch.setattr(
@@ -748,32 +797,34 @@ async def test_user_memory_postprocessor_skips_notice_when_disabled(
         "unique_user_memory.user_memory_postprocessor.upload_user_memory",
         AsyncMock(return_value=True),
     )
-    chat_service = MagicMock()
-    chat_service.modify_assistant_message_async = AsyncMock()
-    monkeypatch.setattr(
-        "unique_user_memory.user_memory_postprocessor.ChatService",
-        MagicMock(return_value=chat_service),
-    )
+    message_step_logger = MagicMock()
+    message_step_logger.log_updating_start = AsyncMock()
+    message_step_logger.log_updating_complete = AsyncMock()
     event = MagicMock()
     event.user_id = "user_1"
     event.company_id = "company_1"
     event.payload.user_message.text = "remember this"
     loop_response = MagicMock()
-    loop_response.message.text = "answer"
-    loop_response.message.id = "msg_1"
-    loop_response.message.references = []
+    loop_response.message.text = "Here is your answer."
     postprocessor = UserMemoryPostprocessor(
-        config=UserMemoryConfig(updating_notice_enabled=False),
+        config=UserMemoryConfig(),
         language_model=_TEST_LANGUAGE_MODEL,
         event=event,
-        state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
+        state=UserMemoryState(
+            scope_id="scope_1",
+            text=empty_profile("user_1"),
+        ),
         logger=MagicMock(),
-        chat_service=chat_service,
+        message_step_logger=message_step_logger,
     )
 
     await postprocessor.run(loop_response)
 
-    chat_service.modify_assistant_message_async.assert_not_awaited()
+    message_step_logger.log_updating_start.assert_awaited_once_with()
+    assert [
+        call.kwargs
+        for call in message_step_logger.log_updating_complete.await_args_list
+    ] == [{"with_settings_entry": False}, {"with_settings_entry": True}]
 
 
 @pytest.mark.asyncio
@@ -786,14 +837,14 @@ async def test_download_user_memory_returns_empty_when_file_missing(
         search_contents,
     )
 
-    result = await download_user_memory(
+    text = await download_user_memory(
         scope_id="scope_1",
         user_id="user_1",
         company_id="company_1",
         logger=MagicMock(),
     )
 
-    assert result == ""
+    assert text == ""
     search_contents.assert_awaited_once_with(
         user_id="user_1",
         company_id="company_1",
@@ -820,14 +871,14 @@ async def test_download_user_memory_downloads_existing_file_to_memory(
         download_content,
     )
 
-    result = await download_user_memory(
+    text = await download_user_memory(
         scope_id="scope_1",
         user_id="user_1",
         company_id="company_1",
         logger=MagicMock(),
     )
 
-    assert result == "# User Memory\n\n## Identity\n- Test"
+    assert text == "# User Memory\n\n## Identity\n- Test"
     search_contents.assert_awaited_once_with(
         user_id="user_1",
         company_id="company_1",
@@ -1013,7 +1064,9 @@ async def test_ensure_user_memory_folder_returns_none_when_access_grant_fails_af
 async def test_upload_user_memory_writes_hidden_skip_ingestion_file(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    upload_content = AsyncMock()
+    uploaded = MagicMock()
+    uploaded.id = "content_uploaded"
+    upload_content = AsyncMock(return_value=uploaded)
     monkeypatch.setattr(
         "unique_user_memory.user_memory.upload_content_from_bytes_async",
         upload_content,
@@ -1063,14 +1116,17 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
     loop_response = MagicMock()
     loop_response.message.text = "noted"
     logger = MagicMock()
-    chat_service = MagicMock()
+    message_step_logger = MagicMock(
+        log_updating_start=AsyncMock(),
+        log_updating_complete=AsyncMock(),
+    )
     postprocessor = UserMemoryPostprocessor(
         config=UserMemoryConfig(),
         language_model=_TEST_LANGUAGE_MODEL,
         event=event,
         state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
         logger=logger,
-        chat_service=chat_service,
+        message_step_logger=message_step_logger,
     )
 
     updated = await postprocessor.run(loop_response)
@@ -1083,11 +1139,117 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
         company_id="company_1",
         logger=logger,
     )
+    message_step_logger.log_updating_complete.assert_awaited_once_with(
+        with_settings_entry=True
+    )
     logger.info.assert_any_call(
         "[user-memory] memory updated and uploaded successfully"
     )
 
 
+@pytest.mark.ai
+@pytest.mark.asyncio
+async def test_user_memory_postprocessor_run_resets_invocation_stats(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Purpose: Verify each run reports only usage attributable to that run.
+    Why this matters: Reused postprocessors must not inflate token analytics.
+    Setup summary: Run twice with distinct usage and assert the second excludes the first.
+    """
+    load_stats = LanguageModelInvocationStats.from_usage(
+        _TEST_LANGUAGE_MODEL.name,
+        LanguageModelTokenUsage(total_tokens=2),
+        source="user_memory_load_condense",
+    )
+    first_run_stats = LanguageModelInvocationStats.from_usage(
+        _TEST_LANGUAGE_MODEL.name,
+        LanguageModelTokenUsage(total_tokens=3),
+        source="user_memory_consolidate_first",
+    )
+    second_run_stats = LanguageModelInvocationStats.from_usage(
+        _TEST_LANGUAGE_MODEL.name,
+        LanguageModelTokenUsage(total_tokens=5),
+        source="user_memory_consolidate_second",
+    )
+    run_stats = iter((first_run_stats, second_run_stats))
+
+    async def consolidate(*, invocation_stats, **kwargs) -> str:  # type: ignore[no-untyped-def]
+        invocation_stats.append(next(run_stats))
+        return "# User Memory\n\n## Identity\n- unchanged"
+
+    monkeypatch.setattr(
+        "unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
+        consolidate,
+    )
+    event = MagicMock()
+    event.user_id = "user_1"
+    event.company_id = "company_1"
+    event.payload.user_message.text = "remember this"
+    loop_response = MagicMock()
+    loop_response.message.text = "noted"
+    state = UserMemoryState(
+        scope_id="scope_1",
+        text="# User Memory\n\n## Identity\n- unchanged",
+        load_invocation_stats=(load_stats,),
+    )
+    postprocessor = UserMemoryPostprocessor(
+        config=UserMemoryConfig(),
+        language_model=_TEST_LANGUAGE_MODEL,
+        event=event,
+        state=state,
+        logger=MagicMock(),
+        message_step_logger=MagicMock(
+            log_updating_start=AsyncMock(),
+            log_updating_complete=AsyncMock(),
+        ),
+    )
+
+    await postprocessor.run(loop_response)
+    first_reported_stats = postprocessor.invocation_stats
+    await postprocessor.run(loop_response)
+
+    assert first_reported_stats == [load_stats, first_run_stats]
+    assert postprocessor.invocation_stats == [second_run_stats]
+
+
+@pytest.mark.ai
+def test_user_memory_postprocessor_take_pending_invocation_stats_drains_once() -> None:
+    """Purpose: Verify load-time usage is reported exactly once, however it's read.
+    Why this matters: A turn that exits before `run()` (cancellation, empty
+    response, a control-taking tool) must still report the load-time condense
+    tokens, and a turn that does reach `run()` must not double-count them.
+    Setup summary: Take the pending stats directly, then run(), and assert
+    run() no longer reports the already-taken load stats.
+    """
+    load_stats = LanguageModelInvocationStats.from_usage(
+        _TEST_LANGUAGE_MODEL.name,
+        LanguageModelTokenUsage(total_tokens=2),
+        source="user_memory_load_condense",
+    )
+    event = MagicMock()
+    event.user_id = "user_1"
+    event.company_id = "company_1"
+    event.payload.user_message.text = "remember this"
+    state = UserMemoryState(
+        scope_id="scope_1",
+        text="# User Memory\n\n## Identity\n- unchanged",
+        load_invocation_stats=(load_stats,),
+    )
+    postprocessor = UserMemoryPostprocessor(
+        config=UserMemoryConfig(),
+        language_model=_TEST_LANGUAGE_MODEL,
+        event=event,
+        state=state,
+        logger=MagicMock(),
+        message_step_logger=MagicMock(),
+    )
+
+    taken = postprocessor.take_pending_invocation_stats()
+
+    assert taken == [load_stats]
+    assert postprocessor.take_pending_invocation_stats() == []
+
+
 @pytest.mark.asyncio
 async def test_user_memory_postprocessor_does_not_log_success_when_upload_fails(
     monkeypatch: pytest.MonkeyPatch,
@@ -1108,14 +1270,16 @@ async def test_user_memory_postprocessor_does_not_log_success_when_upload_fails(
     loop_response = MagicMock()
     loop_response.message.text = "noted"
     logger = MagicMock()
-    chat_service = MagicMock()
     postprocessor = UserMemoryPostprocessor(
         config=UserMemoryConfig(),
         language_model=_TEST_LANGUAGE_MODEL,
         event=event,
         state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
         logger=logger,
-        chat_service=chat_service,
+        message_step_logger=MagicMock(
+            log_updating_start=AsyncMock(),
+            log_updating_complete=AsyncMock(),
+        ),
     )
 
     updated = await postprocessor.run(loop_response)

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev3}/unique_user_memory/user_memory.py

@@ -21,6 +21,9 @@ from unique_toolkit.language_model import (
     TypeEncoder,
 )
 from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.invocation_stats import (
+    LanguageModelInvocationStats,
+)
 
 from unique_user_memory.config import UserMemoryConfig
 from unique_user_memory.user_memory_prompts import (
@@ -97,6 +100,7 @@ def _restore_frontmatter(original: str, body: str) -> str:
 class UserMemoryState:
     scope_id: str
     text: str
+    load_invocation_stats: tuple[LanguageModelInvocationStats, ...] = ()
 
 
 def _get_model_tokenizer(
@@ -181,6 +185,8 @@ async def condense_user_memory(
     language_model: LanguageModelInfo,
     event: ChatEvent,
     logger: Logger,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
+    invocation_source: str = "user_memory_condense",
 ) -> str | None:
     """Ask the LLM to rewrite an oversized profile into a shorter one.
 
@@ -233,6 +239,15 @@ async def condense_user_memory(
         )
         return None
 
+    if invocation_stats is not None and response.usage is not None:
+        invocation_stats.append(
+            LanguageModelInvocationStats.from_usage(
+                language_model.name,
+                response.usage,
+                source=invocation_source,
+            )
+        )
+
     try:
         raw = response.choices[0].message.content or ""
     except Exception as exc:
@@ -268,6 +283,8 @@ async def fit_user_memory(
     language_model: LanguageModelInfo,
     event: ChatEvent,
     logger: Logger,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
+    invocation_source: str = "user_memory_condense",
 ) -> str:
     """Ensure ``content`` fits ``max_tokens``, condensing before cutting.
 
@@ -293,6 +310,8 @@ async def fit_user_memory(
         language_model=language_model,
         event=event,
         logger=logger,
+        invocation_stats=invocation_stats,
+        invocation_source=invocation_source,
     )
     if condensed is not None:
         condensed = _restore_frontmatter(content, condensed)
@@ -358,6 +377,7 @@ async def load_user_memory(
         company_id=company_id,
         logger=logger,
     )
+    invocation_stats: list[LanguageModelInvocationStats] = []
     return UserMemoryState(
         scope_id=scope_id,
         text=await fit_user_memory(
@@ -366,7 +386,10 @@ async def load_user_memory(
             language_model=language_model,
             event=event,
             logger=logger,
+            invocation_stats=invocation_stats,
+            invocation_source="user_memory_load_condense",
         ),
+        load_invocation_stats=tuple(invocation_stats),
     )
 
 
@@ -614,6 +637,7 @@ async def should_consolidate_memory(
     language_model: LanguageModelInfo,
     event: ChatEvent,
     logger: Logger,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
 ) -> bool:
     """Cheaply decide whether the turn warrants a full memory rewrite.
 
@@ -664,6 +688,15 @@ async def should_consolidate_memory(
         )
         return True
 
+    if invocation_stats is not None and response.usage is not None:
+        invocation_stats.append(
+            LanguageModelInvocationStats.from_usage(
+                language_model.name,
+                response.usage,
+                source="user_memory_gate",
+            )
+        )
+
     try:
         raw = response.choices[0].message.content or ""
     except Exception as exc:
@@ -702,6 +735,7 @@ async def consolidate_user_memory(
     logger: Logger,
     on_update_start: Callable[[], Awaitable[None]] = noop_update_callback,
     on_update_end: Callable[[], Awaitable[None]] = noop_update_callback,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
 ) -> str:
     """Consolidate the latest turn into the user's memory profile.
 
@@ -739,6 +773,7 @@ async def consolidate_user_memory(
         language_model=language_model,
         event=event,
         logger=logger,
+        invocation_stats=invocation_stats,
     ):
         return safe_current
 
@@ -753,6 +788,7 @@ async def consolidate_user_memory(
             language_model=language_model,
             event=event,
             logger=logger,
+            invocation_stats=invocation_stats,
         )
     finally:
         await on_update_end()
@@ -768,6 +804,7 @@ async def _rewrite_user_memory(
     language_model: LanguageModelInfo,
     event: ChatEvent,
     logger: Logger,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
 ) -> str:
     if not safe_current.strip():
         safe_current = empty_profile(user_id)
@@ -816,6 +853,15 @@ async def _rewrite_user_memory(
         )
         return safe_current
 
+    if invocation_stats is not None and response.usage is not None:
+        invocation_stats.append(
+            LanguageModelInvocationStats.from_usage(
+                language_model.name,
+                response.usage,
+                source="user_memory_consolidation",
+            )
+        )
+
     try:
         raw = response.choices[0].message.content or ""
     except Exception as exc:
@@ -860,6 +906,8 @@ async def _rewrite_user_memory(
         language_model=language_model,
         event=event,
         logger=logger,
+        invocation_stats=invocation_stats,
+        invocation_source="user_memory_post_consolidation_condense",
     )
     logger.info(
         "[user-memory] consolidation produced %d tokens (cap=%d)",

unique_user_memory-2026.32.0.dev3/unique_user_memory/user_memory_message_log.py

@@ -0,0 +1,139 @@
+"""MessageLog Steps for user-memory load and update."""
+
+from logging import Logger, getLogger
+from typing import Literal
+
+from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
+from unique_toolkit.chat.schemas import (
+    MessageLog,
+    MessageLogDetails,
+    MessageLogEvent,
+    MessageLogStatus,
+)
+
+_LOGGER = getLogger(__name__)
+
+_LOADING_HEADER = "Loading context memory"
+_UPDATING_HEADER = "Updating your memory"
+_CONTEXT_MEMORY_ENTRY_TEXT = "Context memory"
+_REVIEW_MEMORY_ENTRY_TEXT = "Review your context memory"
+
+# Typed MessageLog detail entry recognised by the chat frontend, which renders
+# it as a badge that opens Settings → Context Memory. Frontends that don't know
+# the type parse it as "Unknown" and render nothing, so emitting it is always
+# safe regardless of deploy order.
+USER_MEMORY_EVENT_TYPE: Literal["UserMemory"] = "UserMemory"
+
+
+def _memory_settings_details(*, text: str) -> MessageLogDetails:
+    return MessageLogDetails(
+        data=[MessageLogEvent(type=USER_MEMORY_EVENT_TYPE, text=text)]
+    )
+
+
+class UserMemoryMessageLogger:
+    """Emits Steps for loading and updating the user's context memory file."""
+
+    def __init__(
+        self,
+        message_step_logger: MessageStepLogger,
+        *,
+        logger: Logger | None = None,
+    ) -> None:
+        self._message_step_logger = message_step_logger
+        self._logger = logger or _LOGGER
+        self._loading_log: MessageLog | None = None
+        self._updating_log: MessageLog | None = None
+
+    async def log_loading_start(self) -> None:
+        await self._safe_create_or_update(
+            active_attr="_loading_log",
+            header=_LOADING_HEADER,
+            status=MessageLogStatus.RUNNING,
+            details=None,
+            action="start loading step",
+        )
+
+    async def log_loading_complete(self, *, with_settings_entry: bool = True) -> None:
+        # Attach the settings entry on the loading step itself (same pattern
+        # as update) — a separate MessageLog with empty text is dropped /
+        # invisible in the chat Steps UI.
+        details = (
+            _memory_settings_details(text=_CONTEXT_MEMORY_ENTRY_TEXT)
+            if with_settings_entry
+            else None
+        )
+        await self._safe_create_or_update(
+            active_attr="_loading_log",
+            header=_LOADING_HEADER,
+            status=MessageLogStatus.COMPLETED,
+            details=details,
+            action="complete loading step",
+        )
+
+    async def log_loading_failed(self) -> None:
+        # Always close the RUNNING step when load raises; otherwise the chat
+        # Steps UI leaves "Loading context memory" stuck for that turn.
+        await self._safe_create_or_update(
+            active_attr="_loading_log",
+            header=_LOADING_HEADER,
+            status=MessageLogStatus.FAILED,
+            details=None,
+            action="fail loading step",
+        )
+
+    async def log_updating_start(self) -> None:
+        await self._safe_create_or_update(
+            active_attr="_updating_log",
+            header=_UPDATING_HEADER,
+            status=MessageLogStatus.RUNNING,
+            details=None,
+            action="start updating step",
+        )
+
+    async def log_updating_complete(self, *, with_settings_entry: bool = False) -> None:
+        # Review entry only after memory was actually written (caller sets
+        # with_settings_entry=True post-upload). While consolidating / on
+        # failed upload the step completes without it.
+        details = (
+            _memory_settings_details(text=_REVIEW_MEMORY_ENTRY_TEXT)
+            if with_settings_entry
+            else None
+        )
+        await self._safe_create_or_update(
+            active_attr="_updating_log",
+            header=_UPDATING_HEADER,
+            status=MessageLogStatus.COMPLETED,
+            details=details,
+            action="complete updating step",
+        )
+
+    async def _safe_create_or_update(
+        self,
+        *,
+        active_attr: str,
+        header: str,
+        status: MessageLogStatus,
+        details: MessageLogDetails | None,
+        action: str,
+    ) -> None:
+        try:
+            active = getattr(self, active_attr)
+            updated = (
+                await self._message_step_logger.create_or_update_message_log_async(
+                    active_message_log=active,
+                    header=header,
+                    status=status,
+                    details=details,
+                    references=[],
+                )
+            )
+            if updated is not None:
+                setattr(self, active_attr, updated)
+        except Exception as exc:
+            self._logger.warning(
+                "[user-memory] failed to %s: [%s] %s",
+                action,
+                type(exc).__name__,
+                exc,
+            )

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev3}/unique_user_memory/user_memory_postprocessor.py

@@ -3,25 +3,20 @@ from logging import Logger
 
 from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
 from unique_toolkit.app.schemas import ChatEvent
-from unique_toolkit.chat.service import ChatService
-from unique_toolkit.content.schemas import ContentReference
 from unique_toolkit.language_model.default_language_model import (
     DEFAULT_LANGUAGE_MODEL,
 )
 from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
 from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
 
 from unique_user_memory.config import UserMemoryConfig
 from unique_user_memory.user_memory import (
     UserMemoryState,
     consolidate_user_memory,
-    noop_update_callback,
     upload_user_memory,
 )
-
-# Transient marker appended to the assistant message while the (slow) memory
-# rewrite runs; removed again once consolidation finishes.
-_UPDATING_NOTICE = "\n\n---\n\n🧠 _Updating context memory…_"
+from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
 
 
 class UserMemoryPostprocessor(Postprocessor):
@@ -35,7 +30,7 @@ class UserMemoryPostprocessor(Postprocessor):
         event: ChatEvent,
         state: UserMemoryState,
         logger: Logger,
-        chat_service: ChatService,
+        message_step_logger: UserMemoryMessageLogger,
     ) -> None:
         super().__init__(name="UserMemoryPostprocessor")
         self._config = config
@@ -48,7 +43,29 @@ class UserMemoryPostprocessor(Postprocessor):
         self._state = state
         self._logger = logger
         self._new_memory: str | None = None
-        self._chat_service: ChatService = chat_service
+        self._pending_load_invocation_stats = list(state.load_invocation_stats)
+        self._invocation_stats: list[LanguageModelInvocationStats] = []
+        self._message_step_logger = message_step_logger
+
+    @property
+    def invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        return list(self._invocation_stats)
+
+    def take_pending_invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        """Pop load-time condense stats not yet reported.
+
+        `UniqueAI` calls this unconditionally at the start of every turn so a
+        turn that exits before `run()` (cancellation, empty response, a
+        control-taking tool) still reports the tokens spent condensing the
+        loaded profile. If `run()` does execute, it drains the same pending
+        list itself, so whichever of the two runs first "wins" and the other
+        sees an empty list -- the tokens are never double-counted or lost.
+        """
+        stats, self._pending_load_invocation_stats = (
+            self._pending_load_invocation_stats,
+            [],
+        )
+        return stats
 
     async def run(self, loop_response: LanguageModelStreamResponse) -> bool:
         """Consolidate and upload user memory for this turn.
@@ -56,37 +73,24 @@ class UserMemoryPostprocessor(Postprocessor):
         Returns True if the memory profile changed and was uploaded, False
         otherwise (no user/company, NOOP consolidation, or failed upload).
         """
+        self._invocation_stats = self.take_pending_invocation_stats()
         self._logger.info("[user-memory] running postprocessor")
         user_id = self._event.user_id
         company_id = self._event.company_id
         if not user_id or not company_id:
             return False
 
-        on_update_start: Callable[[], Awaitable[None]] = noop_update_callback
-        on_update_end: Callable[[], Awaitable[None]] = noop_update_callback
-        if self._config.updating_notice_enabled:
-            original_text = loop_response.message.text or ""
-            message_id = loop_response.message.id
-            references = loop_response.message.references
-
-            async def _on_update_start() -> None:
-                await self._set_message_content(
-                    content=original_text + _UPDATING_NOTICE,
-                    message_id=message_id,
-                    references=references,
-                    action="show updating notice",
-                )
-
-            async def _on_update_end() -> None:
-                await self._set_message_content(
-                    content=original_text,
-                    message_id=message_id,
-                    references=references,
-                    action="remove updating notice",
-                )
-
-            on_update_start = _on_update_start
-            on_update_end = _on_update_end
+        async def _on_update_start() -> None:
+            await self._message_step_logger.log_updating_start()
+
+        async def _on_update_end() -> None:
+            # Complete without the review entry; attach it only after upload.
+            await self._message_step_logger.log_updating_complete(
+                with_settings_entry=False
+            )
+
+        on_update_start: Callable[[], Awaitable[None]] = _on_update_start
+        on_update_end: Callable[[], Awaitable[None]] = _on_update_end
 
         self._new_memory = await consolidate_user_memory(
             current_memory=self._state.text,
@@ -99,6 +103,7 @@ class UserMemoryPostprocessor(Postprocessor):
             logger=self._logger,
             on_update_start=on_update_start,
             on_update_end=on_update_end,
+            invocation_stats=self._invocation_stats,
         )
 
         if self._new_memory == self._state.text:
@@ -116,31 +121,10 @@ class UserMemoryPostprocessor(Postprocessor):
             self._logger.warning("[user-memory] memory update was not uploaded")
             return False
 
+        await self._message_step_logger.log_updating_complete(with_settings_entry=True)
         self._logger.info("[user-memory] memory updated and uploaded successfully")
         return True
 
-    async def _set_message_content(
-        self,
-        *,
-        content: str,
-        message_id: str | None,
-        references: list[ContentReference] | None,
-        action: str,
-    ) -> None:
-        try:
-            await self._chat_service.modify_assistant_message_async(
-                content=content,
-                message_id=message_id,
-                references=references,
-            )
-        except Exception as exc:
-            self._logger.warning(
-                "[user-memory] failed to %s: [%s] %s",
-                action,
-                type(exc).__name__,
-                exc,
-            )
-
     def apply_postprocessing_to_response(
         self, loop_response: LanguageModelStreamResponse
     ) -> bool: