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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,14 +1,14 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: unique-user-memory
3
- Version: 2026.32.0.dev2
3
+ Version: 2026.32.0.dev3
4
4
  Summary:
5
5
  Author: Fabian Schläpfer
6
6
  Author-email: Fabian Schläpfer <fabian@unique.ch>
7
7
  License: Proprietary
8
8
  Requires-Dist: jinja2>=3.1.6
9
9
  Requires-Dist: pydantic>=2.8.2
10
- Requires-Dist: unique-sdk>=2026.32.0.dev5,<2026.32.0rc0
11
- Requires-Dist: unique-toolkit>=2026.32.0.dev4,<2026.32.0rc0
10
+ Requires-Dist: unique-sdk>=2026.32.0.dev8,<2026.32.0rc0
11
+ Requires-Dist: unique-toolkit>=2026.32.0.dev8,<2026.32.0rc0
12
12
  Requires-Python: >=3.12, <4
13
13
  Description-Content-Type: text/markdown
14
14
 
@@ -23,7 +23,8 @@ Persistent per-user memory for Unique AI agents.
23
23
  The package provides:
24
24
 
25
25
  - `UserMemoryConfig` - Pydantic configuration for the consolidation model, profile token budget, and memory folder.
26
- - `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).
26
+ - `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.
27
+ - `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.
27
28
  - `UserMemoryPostprocessor` - runs after the assistant response, consolidates the latest turn into the profile, and uploads the updated `memory.md`.
28
29
 
29
30
  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
31
32
  ## Lifecycle
32
33
 
33
34
  1. The orchestrator enables memory when `space.allow_user_memory` is true.
34
- 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.
35
- 3. The loaded memory text is passed into the agent context for the current turn.
36
- 4. `UserMemoryPostprocessor` runs after the assistant response.
37
- 5. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
38
- 6. If the profile changed, `memory.md` is uploaded back to the user's folder with ingestion skipped and the content hidden from chat.
35
+ 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.
36
+ 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.
37
+ 4. If memory was loaded, its text is passed into the agent context for the current turn.
38
+ 5. `UserMemoryPostprocessor` runs after the assistant response.
39
+ 6. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
40
+ 7. If a rewrite runs, an **Updating your memory** Step is shown while consolidating (no settings entry yet).
41
+ 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.
39
42
 
40
43
  ## Storage Model
41
44
 
@@ -109,7 +112,9 @@ Typical orchestration code loads memory before the agent loop and registers the
109
112
  `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:
110
113
 
111
114
  ```python
115
+ from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
112
116
  from unique_user_memory.user_memory import load_user_memory
117
+ from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
113
118
  from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
114
119
 
115
120
  user_memory_config = config.agent.services.user_memory_config
@@ -122,14 +127,36 @@ memory_language_model = (
122
127
  else user_memory_config.language_model
123
128
  )
124
129
 
125
- user_memory_state = await load_user_memory(
126
- event=event,
127
- config=user_memory_config,
128
- language_model=memory_language_model,
130
+ message_step_logger = MessageStepLogger(chat_service)
131
+ memory_message_step_logger = UserMemoryMessageLogger(
132
+ message_step_logger,
129
133
  logger=logger,
130
134
  )
131
-
132
- if user_memory_state is not None:
135
+ await memory_message_step_logger.log_loading_start()
136
+ user_memory_state = None
137
+ load_succeeded = False
138
+ try:
139
+ user_memory_state = await load_user_memory(
140
+ event=event,
141
+ config=user_memory_config,
142
+ language_model=memory_language_model,
143
+ logger=logger,
144
+ )
145
+ load_succeeded = True
146
+ except Exception as exc:
147
+ logger.warning(
148
+ "[user-memory] load raised - running without memory: [%s] %s",
149
+ type(exc).__name__,
150
+ exc,
151
+ )
152
+ finally:
153
+ # Always close the RUNNING step — otherwise the chat Steps UI stays stuck
154
+ # on "Loading context memory" for that turn when load raises.
155
+ if not load_succeeded:
156
+ await memory_message_step_logger.log_loading_failed()
157
+
158
+ if load_succeeded and user_memory_state is not None:
159
+ await memory_message_step_logger.log_loading_complete(with_settings_entry=True)
133
160
  user_memory_text = user_memory_state.text
134
161
  postprocessor_manager.add_postprocessor(
135
162
  UserMemoryPostprocessor(
@@ -138,8 +165,11 @@ if user_memory_state is not None:
138
165
  event=event,
139
166
  state=user_memory_state,
140
167
  logger=logger,
168
+ message_step_logger=memory_message_step_logger,
141
169
  )
142
170
  )
171
+ elif load_succeeded:
172
+ await memory_message_step_logger.log_loading_complete(with_settings_entry=False)
143
173
  ```
144
174
 
145
175
  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.
@@ -9,7 +9,8 @@ Persistent per-user memory for Unique AI agents.
9
9
  The package provides:
10
10
 
11
11
  - `UserMemoryConfig` - Pydantic configuration for the consolidation model, profile token budget, and memory folder.
12
- - `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).
12
+ - `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.
13
+ - `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.
13
14
  - `UserMemoryPostprocessor` - runs after the assistant response, consolidates the latest turn into the profile, and uploads the updated `memory.md`.
14
15
 
15
16
  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
17
18
  ## Lifecycle
18
19
 
19
20
  1. The orchestrator enables memory when `space.allow_user_memory` is true.
20
- 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.
21
- 3. The loaded memory text is passed into the agent context for the current turn.
22
- 4. `UserMemoryPostprocessor` runs after the assistant response.
23
- 5. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
24
- 6. If the profile changed, `memory.md` is uploaded back to the user's folder with ingestion skipped and the content hidden from chat.
21
+ 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.
22
+ 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.
23
+ 4. If memory was loaded, its text is passed into the agent context for the current turn.
24
+ 5. `UserMemoryPostprocessor` runs after the assistant response.
25
+ 6. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
26
+ 7. If a rewrite runs, an **Updating your memory** Step is shown while consolidating (no settings entry yet).
27
+ 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.
25
28
 
26
29
  ## Storage Model
27
30
 
@@ -95,7 +98,9 @@ Typical orchestration code loads memory before the agent loop and registers the
95
98
  `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:
96
99
 
97
100
  ```python
101
+ from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
98
102
  from unique_user_memory.user_memory import load_user_memory
103
+ from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
99
104
  from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
100
105
 
101
106
  user_memory_config = config.agent.services.user_memory_config
@@ -108,14 +113,36 @@ memory_language_model = (
108
113
  else user_memory_config.language_model
109
114
  )
110
115
 
111
- user_memory_state = await load_user_memory(
112
- event=event,
113
- config=user_memory_config,
114
- language_model=memory_language_model,
116
+ message_step_logger = MessageStepLogger(chat_service)
117
+ memory_message_step_logger = UserMemoryMessageLogger(
118
+ message_step_logger,
115
119
  logger=logger,
116
120
  )
117
-
118
- if user_memory_state is not None:
121
+ await memory_message_step_logger.log_loading_start()
122
+ user_memory_state = None
123
+ load_succeeded = False
124
+ try:
125
+ user_memory_state = await load_user_memory(
126
+ event=event,
127
+ config=user_memory_config,
128
+ language_model=memory_language_model,
129
+ logger=logger,
130
+ )
131
+ load_succeeded = True
132
+ except Exception as exc:
133
+ logger.warning(
134
+ "[user-memory] load raised - running without memory: [%s] %s",
135
+ type(exc).__name__,
136
+ exc,
137
+ )
138
+ finally:
139
+ # Always close the RUNNING step — otherwise the chat Steps UI stays stuck
140
+ # on "Loading context memory" for that turn when load raises.
141
+ if not load_succeeded:
142
+ await memory_message_step_logger.log_loading_failed()
143
+
144
+ if load_succeeded and user_memory_state is not None:
145
+ await memory_message_step_logger.log_loading_complete(with_settings_entry=True)
119
146
  user_memory_text = user_memory_state.text
120
147
  postprocessor_manager.add_postprocessor(
121
148
  UserMemoryPostprocessor(
@@ -124,8 +151,11 @@ if user_memory_state is not None:
124
151
  event=event,
125
152
  state=user_memory_state,
126
153
  logger=logger,
154
+ message_step_logger=memory_message_step_logger,
127
155
  )
128
156
  )
157
+ elif load_succeeded:
158
+ await memory_message_step_logger.log_loading_complete(with_settings_entry=False)
129
159
  ```
130
160
 
131
161
  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.
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "unique_user_memory"
3
- version = "2026.32.0.dev2"
3
+ version = "2026.32.0.dev3"
4
4
  description = ""
5
5
  authors = [
6
6
  { name = "Fabian Schläpfer", email = "fabian@unique.ch" },
@@ -11,8 +11,8 @@ requires-python = ">=3.12,<4"
11
11
  dependencies = [
12
12
  "jinja2>=3.1.6",
13
13
  "pydantic>=2.8.2",
14
- "unique-sdk>=2026.32.0.dev5,<2026.32.0rc0",
15
- "unique-toolkit>=2026.32.0.dev4,<2026.32.0rc0",
14
+ "unique-sdk>=2026.32.0.dev8,<2026.32.0rc0",
15
+ "unique-toolkit>=2026.32.0.dev8,<2026.32.0rc0",
16
16
  ]
17
17
 
18
18
  [dependency-groups]
@@ -42,14 +42,6 @@ class UserMemoryConfig(BaseModel):
42
42
  "runs."
43
43
  ),
44
44
  )
45
- updating_notice_enabled: bool = Field(
46
- default=True,
47
- description=(
48
- "When true, a transient 'updating context memory' notice is appended to the "
49
- "assistant message while the memory rewrite runs, and removed "
50
- "again once it completes."
51
- ),
52
- )
53
45
  root_folder: Annotated[str, RJSFMetaTag.SpecialWidget.hidden()] = Field(
54
46
  default="user-memory",
55
47
  min_length=1,
@@ -2,6 +2,7 @@ import asyncio
2
2
  from unittest.mock import AsyncMock, MagicMock
3
3
 
4
4
  import pytest
5
+ from unique_toolkit.chat.schemas import MessageLogStatus
5
6
  from unique_toolkit.language_model.default_language_model import (
6
7
  DEFAULT_LANGUAGE_MODEL,
7
8
  )
@@ -20,10 +21,10 @@ from unique_user_memory.user_memory import (
20
21
  enforce_token_cap,
21
22
  ensure_user_memory_folder,
22
23
  fit_user_memory,
23
- noop_update_callback,
24
24
  should_consolidate_memory,
25
25
  upload_user_memory,
26
26
  )
27
+ from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
27
28
  from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
28
29
  from unique_user_memory.user_memory_prompts import (
29
30
  consolidation_system_prompt,
@@ -679,67 +680,113 @@ async def test_consolidate_user_memory_invokes_update_end_when_start_cancelled(
679
680
 
680
681
 
681
682
  @pytest.mark.asyncio
682
- async def test_user_memory_postprocessor_shows_and_removes_updating_notice(
683
- monkeypatch: pytest.MonkeyPatch,
684
- ) -> None:
685
- updated_memory = "# User Memory\n\n## Identity\n- Updated"
686
- original_text = "Here is your answer."
683
+ async def test_user_memory_message_step_logger_load_emits_settings_detail_entry() -> (
684
+ None
685
+ ):
686
+ step_logger = MagicMock()
687
+ step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
688
+ message_logger = UserMemoryMessageLogger(step_logger)
689
+
690
+ await message_logger.log_loading_start()
691
+ await message_logger.log_loading_complete(with_settings_entry=True)
692
+
693
+ assert step_logger.create_or_update_message_log_async.await_count == 2
694
+ start_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
695
+ 0
696
+ ].kwargs
697
+ assert start_kwargs["header"] == "Loading context memory"
698
+ assert start_kwargs["details"] is None
699
+ assert start_kwargs["references"] == []
700
+ complete_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
701
+ 1
702
+ ].kwargs
703
+ assert complete_kwargs["status"] == MessageLogStatus.COMPLETED
704
+ assert complete_kwargs["references"] == []
705
+ entry = complete_kwargs["details"].data[0]
706
+ assert entry.type == "UserMemory"
707
+ assert entry.text == "Context memory"
687
708
 
688
- async def fake_consolidate(*, on_update_start, on_update_end, **kwargs) -> str: # type: ignore[no-untyped-def]
689
- await on_update_start()
690
- await on_update_end()
691
- return updated_memory
692
709
 
693
- monkeypatch.setattr(
694
- "unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
695
- fake_consolidate,
696
- )
697
- monkeypatch.setattr(
698
- "unique_user_memory.user_memory_postprocessor.upload_user_memory",
699
- AsyncMock(return_value=True),
700
- )
701
- chat_service = MagicMock()
702
- chat_service.modify_assistant_message_async = AsyncMock()
703
- monkeypatch.setattr(
704
- "unique_user_memory.user_memory_postprocessor.ChatService",
705
- MagicMock(return_value=chat_service),
706
- )
707
- event = MagicMock()
708
- event.user_id = "user_1"
709
- event.company_id = "company_1"
710
- event.payload.user_message.text = "remember this"
711
- loop_response = MagicMock()
712
- loop_response.message.text = original_text
713
- loop_response.message.id = "msg_1"
714
- loop_response.message.references = []
715
- postprocessor = UserMemoryPostprocessor(
716
- config=UserMemoryConfig(),
717
- language_model=_TEST_LANGUAGE_MODEL,
718
- event=event,
719
- state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
720
- logger=MagicMock(),
721
- chat_service=chat_service,
722
- )
710
+ @pytest.mark.asyncio
711
+ async def test_user_memory_message_step_logger_load_skips_entry_when_disabled() -> None:
712
+ step_logger = MagicMock()
713
+ step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
714
+ message_logger = UserMemoryMessageLogger(step_logger)
723
715
 
724
- await postprocessor.run(loop_response)
716
+ await message_logger.log_loading_complete(with_settings_entry=False)
717
+
718
+ complete_kwargs = step_logger.create_or_update_message_log_async.await_args.kwargs
719
+ assert complete_kwargs["details"] is None
720
+ assert complete_kwargs["references"] == []
721
+
722
+
723
+ @pytest.mark.asyncio
724
+ async def test_user_memory_message_step_logger_load_failed_marks_failed_status() -> (
725
+ None
726
+ ):
727
+ """
728
+ Purpose: log_loading_failed updates the loading Step to FAILED with no
729
+ settings entry.
730
+ Why this matters: Callers must close a RUNNING loading Step when load
731
+ raises so the chat UI does not stay stuck.
732
+ Setup summary: Start then fail the loading step; assert FAILED status.
733
+ """
734
+ step_logger = MagicMock()
735
+ step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
736
+ message_logger = UserMemoryMessageLogger(step_logger)
725
737
 
726
- calls = chat_service.modify_assistant_message_async.await_args_list
727
- assert len(calls) == 2
728
- assert calls[0].kwargs["content"].startswith(original_text)
729
- assert calls[0].kwargs["content"] != original_text
730
- assert calls[0].kwargs["message_id"] == "msg_1"
731
- assert calls[1].kwargs["content"] == original_text
738
+ await message_logger.log_loading_start()
739
+ await message_logger.log_loading_failed()
740
+
741
+ failed_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
742
+ -1
743
+ ].kwargs
744
+ assert failed_kwargs["header"] == "Loading context memory"
745
+ assert failed_kwargs["status"] == MessageLogStatus.FAILED
746
+ assert failed_kwargs["details"] is None
747
+ assert failed_kwargs["references"] == []
748
+
749
+
750
+ @pytest.mark.asyncio
751
+ async def test_user_memory_message_step_logger_update_attaches_review_entry_only_when_requested() -> (
752
+ None
753
+ ):
754
+ step_logger = MagicMock()
755
+ step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
756
+ message_logger = UserMemoryMessageLogger(step_logger)
757
+
758
+ await message_logger.log_updating_start()
759
+ await message_logger.log_updating_complete(with_settings_entry=False)
760
+ await message_logger.log_updating_complete(with_settings_entry=True)
761
+
762
+ assert step_logger.create_or_update_message_log_async.await_count == 3
763
+ start_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
764
+ 0
765
+ ].kwargs
766
+ assert start_kwargs["header"] == "Updating your memory"
767
+ assert start_kwargs["details"] is None
768
+ complete_without_entry = (
769
+ step_logger.create_or_update_message_log_async.await_args_list[1].kwargs
770
+ )
771
+ assert complete_without_entry["details"] is None
772
+ complete_with_entry = (
773
+ step_logger.create_or_update_message_log_async.await_args_list[2].kwargs
774
+ )
775
+ assert complete_with_entry["references"] == []
776
+ entry = complete_with_entry["details"].data[0]
777
+ assert entry.type == "UserMemory"
778
+ assert entry.text == "Review your context memory"
732
779
 
733
780
 
734
781
  @pytest.mark.asyncio
735
- async def test_user_memory_postprocessor_skips_notice_when_disabled(
782
+ async def test_user_memory_postprocessor_emits_updating_message_logs(
736
783
  monkeypatch: pytest.MonkeyPatch,
737
784
  ) -> None:
738
785
  updated_memory = "# User Memory\n\n## Identity\n- Updated"
739
786
 
740
787
  async def fake_consolidate(*, on_update_start, on_update_end, **kwargs) -> str: # type: ignore[no-untyped-def]
741
- assert on_update_start is noop_update_callback
742
- assert on_update_end is noop_update_callback
788
+ await on_update_start()
789
+ await on_update_end()
743
790
  return updated_memory
744
791
 
745
792
  monkeypatch.setattr(
@@ -750,32 +797,34 @@ async def test_user_memory_postprocessor_skips_notice_when_disabled(
750
797
  "unique_user_memory.user_memory_postprocessor.upload_user_memory",
751
798
  AsyncMock(return_value=True),
752
799
  )
753
- chat_service = MagicMock()
754
- chat_service.modify_assistant_message_async = AsyncMock()
755
- monkeypatch.setattr(
756
- "unique_user_memory.user_memory_postprocessor.ChatService",
757
- MagicMock(return_value=chat_service),
758
- )
800
+ message_step_logger = MagicMock()
801
+ message_step_logger.log_updating_start = AsyncMock()
802
+ message_step_logger.log_updating_complete = AsyncMock()
759
803
  event = MagicMock()
760
804
  event.user_id = "user_1"
761
805
  event.company_id = "company_1"
762
806
  event.payload.user_message.text = "remember this"
763
807
  loop_response = MagicMock()
764
- loop_response.message.text = "answer"
765
- loop_response.message.id = "msg_1"
766
- loop_response.message.references = []
808
+ loop_response.message.text = "Here is your answer."
767
809
  postprocessor = UserMemoryPostprocessor(
768
- config=UserMemoryConfig(updating_notice_enabled=False),
810
+ config=UserMemoryConfig(),
769
811
  language_model=_TEST_LANGUAGE_MODEL,
770
812
  event=event,
771
- state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
813
+ state=UserMemoryState(
814
+ scope_id="scope_1",
815
+ text=empty_profile("user_1"),
816
+ ),
772
817
  logger=MagicMock(),
773
- chat_service=chat_service,
818
+ message_step_logger=message_step_logger,
774
819
  )
775
820
 
776
821
  await postprocessor.run(loop_response)
777
822
 
778
- chat_service.modify_assistant_message_async.assert_not_awaited()
823
+ message_step_logger.log_updating_start.assert_awaited_once_with()
824
+ assert [
825
+ call.kwargs
826
+ for call in message_step_logger.log_updating_complete.await_args_list
827
+ ] == [{"with_settings_entry": False}, {"with_settings_entry": True}]
779
828
 
780
829
 
781
830
  @pytest.mark.asyncio
@@ -788,14 +837,14 @@ async def test_download_user_memory_returns_empty_when_file_missing(
788
837
  search_contents,
789
838
  )
790
839
 
791
- result = await download_user_memory(
840
+ text = await download_user_memory(
792
841
  scope_id="scope_1",
793
842
  user_id="user_1",
794
843
  company_id="company_1",
795
844
  logger=MagicMock(),
796
845
  )
797
846
 
798
- assert result == ""
847
+ assert text == ""
799
848
  search_contents.assert_awaited_once_with(
800
849
  user_id="user_1",
801
850
  company_id="company_1",
@@ -822,14 +871,14 @@ async def test_download_user_memory_downloads_existing_file_to_memory(
822
871
  download_content,
823
872
  )
824
873
 
825
- result = await download_user_memory(
874
+ text = await download_user_memory(
826
875
  scope_id="scope_1",
827
876
  user_id="user_1",
828
877
  company_id="company_1",
829
878
  logger=MagicMock(),
830
879
  )
831
880
 
832
- assert result == "# User Memory\n\n## Identity\n- Test"
881
+ assert text == "# User Memory\n\n## Identity\n- Test"
833
882
  search_contents.assert_awaited_once_with(
834
883
  user_id="user_1",
835
884
  company_id="company_1",
@@ -1015,7 +1064,9 @@ async def test_ensure_user_memory_folder_returns_none_when_access_grant_fails_af
1015
1064
  async def test_upload_user_memory_writes_hidden_skip_ingestion_file(
1016
1065
  monkeypatch: pytest.MonkeyPatch,
1017
1066
  ) -> None:
1018
- upload_content = AsyncMock()
1067
+ uploaded = MagicMock()
1068
+ uploaded.id = "content_uploaded"
1069
+ upload_content = AsyncMock(return_value=uploaded)
1019
1070
  monkeypatch.setattr(
1020
1071
  "unique_user_memory.user_memory.upload_content_from_bytes_async",
1021
1072
  upload_content,
@@ -1065,14 +1116,17 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
1065
1116
  loop_response = MagicMock()
1066
1117
  loop_response.message.text = "noted"
1067
1118
  logger = MagicMock()
1068
- chat_service = MagicMock()
1119
+ message_step_logger = MagicMock(
1120
+ log_updating_start=AsyncMock(),
1121
+ log_updating_complete=AsyncMock(),
1122
+ )
1069
1123
  postprocessor = UserMemoryPostprocessor(
1070
1124
  config=UserMemoryConfig(),
1071
1125
  language_model=_TEST_LANGUAGE_MODEL,
1072
1126
  event=event,
1073
1127
  state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
1074
1128
  logger=logger,
1075
- chat_service=chat_service,
1129
+ message_step_logger=message_step_logger,
1076
1130
  )
1077
1131
 
1078
1132
  updated = await postprocessor.run(loop_response)
@@ -1085,6 +1139,9 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
1085
1139
  company_id="company_1",
1086
1140
  logger=logger,
1087
1141
  )
1142
+ message_step_logger.log_updating_complete.assert_awaited_once_with(
1143
+ with_settings_entry=True
1144
+ )
1088
1145
  logger.info.assert_any_call(
1089
1146
  "[user-memory] memory updated and uploaded successfully"
1090
1147
  )
@@ -1141,7 +1198,10 @@ async def test_user_memory_postprocessor_run_resets_invocation_stats(
1141
1198
  event=event,
1142
1199
  state=state,
1143
1200
  logger=MagicMock(),
1144
- chat_service=MagicMock(),
1201
+ message_step_logger=MagicMock(
1202
+ log_updating_start=AsyncMock(),
1203
+ log_updating_complete=AsyncMock(),
1204
+ ),
1145
1205
  )
1146
1206
 
1147
1207
  await postprocessor.run(loop_response)
@@ -1181,7 +1241,7 @@ def test_user_memory_postprocessor_take_pending_invocation_stats_drains_once() -
1181
1241
  event=event,
1182
1242
  state=state,
1183
1243
  logger=MagicMock(),
1184
- chat_service=MagicMock(),
1244
+ message_step_logger=MagicMock(),
1185
1245
  )
1186
1246
 
1187
1247
  taken = postprocessor.take_pending_invocation_stats()
@@ -1210,14 +1270,16 @@ async def test_user_memory_postprocessor_does_not_log_success_when_upload_fails(
1210
1270
  loop_response = MagicMock()
1211
1271
  loop_response.message.text = "noted"
1212
1272
  logger = MagicMock()
1213
- chat_service = MagicMock()
1214
1273
  postprocessor = UserMemoryPostprocessor(
1215
1274
  config=UserMemoryConfig(),
1216
1275
  language_model=_TEST_LANGUAGE_MODEL,
1217
1276
  event=event,
1218
1277
  state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
1219
1278
  logger=logger,
1220
- chat_service=chat_service,
1279
+ message_step_logger=MagicMock(
1280
+ log_updating_start=AsyncMock(),
1281
+ log_updating_complete=AsyncMock(),
1282
+ ),
1221
1283
  )
1222
1284
 
1223
1285
  updated = await postprocessor.run(loop_response)
@@ -0,0 +1,139 @@
1
+ """MessageLog Steps for user-memory load and update."""
2
+
3
+ from logging import Logger, getLogger
4
+ from typing import Literal
5
+
6
+ from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
7
+ from unique_toolkit.chat.schemas import (
8
+ MessageLog,
9
+ MessageLogDetails,
10
+ MessageLogEvent,
11
+ MessageLogStatus,
12
+ )
13
+
14
+ _LOGGER = getLogger(__name__)
15
+
16
+ _LOADING_HEADER = "Loading context memory"
17
+ _UPDATING_HEADER = "Updating your memory"
18
+ _CONTEXT_MEMORY_ENTRY_TEXT = "Context memory"
19
+ _REVIEW_MEMORY_ENTRY_TEXT = "Review your context memory"
20
+
21
+ # Typed MessageLog detail entry recognised by the chat frontend, which renders
22
+ # it as a badge that opens Settings → Context Memory. Frontends that don't know
23
+ # the type parse it as "Unknown" and render nothing, so emitting it is always
24
+ # safe regardless of deploy order.
25
+ USER_MEMORY_EVENT_TYPE: Literal["UserMemory"] = "UserMemory"
26
+
27
+
28
+ def _memory_settings_details(*, text: str) -> MessageLogDetails:
29
+ return MessageLogDetails(
30
+ data=[MessageLogEvent(type=USER_MEMORY_EVENT_TYPE, text=text)]
31
+ )
32
+
33
+
34
+ class UserMemoryMessageLogger:
35
+ """Emits Steps for loading and updating the user's context memory file."""
36
+
37
+ def __init__(
38
+ self,
39
+ message_step_logger: MessageStepLogger,
40
+ *,
41
+ logger: Logger | None = None,
42
+ ) -> None:
43
+ self._message_step_logger = message_step_logger
44
+ self._logger = logger or _LOGGER
45
+ self._loading_log: MessageLog | None = None
46
+ self._updating_log: MessageLog | None = None
47
+
48
+ async def log_loading_start(self) -> None:
49
+ await self._safe_create_or_update(
50
+ active_attr="_loading_log",
51
+ header=_LOADING_HEADER,
52
+ status=MessageLogStatus.RUNNING,
53
+ details=None,
54
+ action="start loading step",
55
+ )
56
+
57
+ async def log_loading_complete(self, *, with_settings_entry: bool = True) -> None:
58
+ # Attach the settings entry on the loading step itself (same pattern
59
+ # as update) — a separate MessageLog with empty text is dropped /
60
+ # invisible in the chat Steps UI.
61
+ details = (
62
+ _memory_settings_details(text=_CONTEXT_MEMORY_ENTRY_TEXT)
63
+ if with_settings_entry
64
+ else None
65
+ )
66
+ await self._safe_create_or_update(
67
+ active_attr="_loading_log",
68
+ header=_LOADING_HEADER,
69
+ status=MessageLogStatus.COMPLETED,
70
+ details=details,
71
+ action="complete loading step",
72
+ )
73
+
74
+ async def log_loading_failed(self) -> None:
75
+ # Always close the RUNNING step when load raises; otherwise the chat
76
+ # Steps UI leaves "Loading context memory" stuck for that turn.
77
+ await self._safe_create_or_update(
78
+ active_attr="_loading_log",
79
+ header=_LOADING_HEADER,
80
+ status=MessageLogStatus.FAILED,
81
+ details=None,
82
+ action="fail loading step",
83
+ )
84
+
85
+ async def log_updating_start(self) -> None:
86
+ await self._safe_create_or_update(
87
+ active_attr="_updating_log",
88
+ header=_UPDATING_HEADER,
89
+ status=MessageLogStatus.RUNNING,
90
+ details=None,
91
+ action="start updating step",
92
+ )
93
+
94
+ async def log_updating_complete(self, *, with_settings_entry: bool = False) -> None:
95
+ # Review entry only after memory was actually written (caller sets
96
+ # with_settings_entry=True post-upload). While consolidating / on
97
+ # failed upload the step completes without it.
98
+ details = (
99
+ _memory_settings_details(text=_REVIEW_MEMORY_ENTRY_TEXT)
100
+ if with_settings_entry
101
+ else None
102
+ )
103
+ await self._safe_create_or_update(
104
+ active_attr="_updating_log",
105
+ header=_UPDATING_HEADER,
106
+ status=MessageLogStatus.COMPLETED,
107
+ details=details,
108
+ action="complete updating step",
109
+ )
110
+
111
+ async def _safe_create_or_update(
112
+ self,
113
+ *,
114
+ active_attr: str,
115
+ header: str,
116
+ status: MessageLogStatus,
117
+ details: MessageLogDetails | None,
118
+ action: str,
119
+ ) -> None:
120
+ try:
121
+ active = getattr(self, active_attr)
122
+ updated = (
123
+ await self._message_step_logger.create_or_update_message_log_async(
124
+ active_message_log=active,
125
+ header=header,
126
+ status=status,
127
+ details=details,
128
+ references=[],
129
+ )
130
+ )
131
+ if updated is not None:
132
+ setattr(self, active_attr, updated)
133
+ except Exception as exc:
134
+ self._logger.warning(
135
+ "[user-memory] failed to %s: [%s] %s",
136
+ action,
137
+ type(exc).__name__,
138
+ exc,
139
+ )
@@ -3,8 +3,6 @@ from logging import Logger
3
3
 
4
4
  from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
5
5
  from unique_toolkit.app.schemas import ChatEvent
6
- from unique_toolkit.chat.service import ChatService
7
- from unique_toolkit.content.schemas import ContentReference
8
6
  from unique_toolkit.language_model.default_language_model import (
9
7
  DEFAULT_LANGUAGE_MODEL,
10
8
  )
@@ -16,13 +14,9 @@ from unique_user_memory.config import UserMemoryConfig
16
14
  from unique_user_memory.user_memory import (
17
15
  UserMemoryState,
18
16
  consolidate_user_memory,
19
- noop_update_callback,
20
17
  upload_user_memory,
21
18
  )
22
-
23
- # Transient marker appended to the assistant message while the (slow) memory
24
- # rewrite runs; removed again once consolidation finishes.
25
- _UPDATING_NOTICE = "\n\n---\n\n🧠 _Updating context memory…_"
19
+ from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
26
20
 
27
21
 
28
22
  class UserMemoryPostprocessor(Postprocessor):
@@ -36,7 +30,7 @@ class UserMemoryPostprocessor(Postprocessor):
36
30
  event: ChatEvent,
37
31
  state: UserMemoryState,
38
32
  logger: Logger,
39
- chat_service: ChatService,
33
+ message_step_logger: UserMemoryMessageLogger,
40
34
  ) -> None:
41
35
  super().__init__(name="UserMemoryPostprocessor")
42
36
  self._config = config
@@ -49,9 +43,9 @@ class UserMemoryPostprocessor(Postprocessor):
49
43
  self._state = state
50
44
  self._logger = logger
51
45
  self._new_memory: str | None = None
52
- self._chat_service: ChatService = chat_service
53
46
  self._pending_load_invocation_stats = list(state.load_invocation_stats)
54
47
  self._invocation_stats: list[LanguageModelInvocationStats] = []
48
+ self._message_step_logger = message_step_logger
55
49
 
56
50
  @property
57
51
  def invocation_stats(self) -> list[LanguageModelInvocationStats]:
@@ -86,31 +80,17 @@ class UserMemoryPostprocessor(Postprocessor):
86
80
  if not user_id or not company_id:
87
81
  return False
88
82
 
89
- on_update_start: Callable[[], Awaitable[None]] = noop_update_callback
90
- on_update_end: Callable[[], Awaitable[None]] = noop_update_callback
91
- if self._config.updating_notice_enabled:
92
- original_text = loop_response.message.text or ""
93
- message_id = loop_response.message.id
94
- references = loop_response.message.references
95
-
96
- async def _on_update_start() -> None:
97
- await self._set_message_content(
98
- content=original_text + _UPDATING_NOTICE,
99
- message_id=message_id,
100
- references=references,
101
- action="show updating notice",
102
- )
103
-
104
- async def _on_update_end() -> None:
105
- await self._set_message_content(
106
- content=original_text,
107
- message_id=message_id,
108
- references=references,
109
- action="remove updating notice",
110
- )
111
-
112
- on_update_start = _on_update_start
113
- on_update_end = _on_update_end
83
+ async def _on_update_start() -> None:
84
+ await self._message_step_logger.log_updating_start()
85
+
86
+ async def _on_update_end() -> None:
87
+ # Complete without the review entry; attach it only after upload.
88
+ await self._message_step_logger.log_updating_complete(
89
+ with_settings_entry=False
90
+ )
91
+
92
+ on_update_start: Callable[[], Awaitable[None]] = _on_update_start
93
+ on_update_end: Callable[[], Awaitable[None]] = _on_update_end
114
94
 
115
95
  self._new_memory = await consolidate_user_memory(
116
96
  current_memory=self._state.text,
@@ -141,31 +121,10 @@ class UserMemoryPostprocessor(Postprocessor):
141
121
  self._logger.warning("[user-memory] memory update was not uploaded")
142
122
  return False
143
123
 
124
+ await self._message_step_logger.log_updating_complete(with_settings_entry=True)
144
125
  self._logger.info("[user-memory] memory updated and uploaded successfully")
145
126
  return True
146
127
 
147
- async def _set_message_content(
148
- self,
149
- *,
150
- content: str,
151
- message_id: str | None,
152
- references: list[ContentReference] | None,
153
- action: str,
154
- ) -> None:
155
- try:
156
- await self._chat_service.modify_assistant_message_async(
157
- content=content,
158
- message_id=message_id,
159
- references=references,
160
- )
161
- except Exception as exc:
162
- self._logger.warning(
163
- "[user-memory] failed to %s: [%s] %s",
164
- action,
165
- type(exc).__name__,
166
- exc,
167
- )
168
-
169
128
  def apply_postprocessing_to_response(
170
129
  self, loop_response: LanguageModelStreamResponse
171
130
  ) -> bool: