unique-user-memory 2026.32.0.dev2__tar.gz → 2026.32.0.dev4__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.dev4
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,9 @@ 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
+ - `profile_body(...)` - strips the YAML frontmatter and returns only the Markdown body. Use it whenever the profile is shown to a model; the frontmatter is bookkeeping for consolidation.
28
+ - `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
29
  - `UserMemoryPostprocessor` - runs after the assistant response, consolidates the latest turn into the profile, and uploads the updated `memory.md`.
28
30
 
29
31
  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 +33,13 @@ The memory file is intentionally small and structured. It is rewritten as a full
31
33
  ## Lifecycle
32
34
 
33
35
  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.
36
+ 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.
37
+ 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.
38
+ 4. If memory was loaded, `profile_body(...)` of its text is passed into the agent context for the current turn — the prompt only gets the Markdown body, while the postprocessor keeps the full file because it needs the frontmatter to carry `turn_count` forward.
39
+ 5. `UserMemoryPostprocessor` runs after the assistant response.
40
+ 6. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
41
+ 7. If a rewrite runs, an **Updating your memory** Step is shown while consolidating (no settings entry yet).
42
+ 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
43
 
40
44
  ## Storage Model
41
45
 
@@ -109,7 +113,9 @@ Typical orchestration code loads memory before the agent loop and registers the
109
113
  `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
114
 
111
115
  ```python
112
- from unique_user_memory.user_memory import load_user_memory
116
+ from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
117
+ from unique_user_memory.user_memory import load_user_memory, profile_body
118
+ from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
113
119
  from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
114
120
 
115
121
  user_memory_config = config.agent.services.user_memory_config
@@ -122,15 +128,39 @@ memory_language_model = (
122
128
  else user_memory_config.language_model
123
129
  )
124
130
 
125
- user_memory_state = await load_user_memory(
126
- event=event,
127
- config=user_memory_config,
128
- language_model=memory_language_model,
131
+ message_step_logger = MessageStepLogger(chat_service)
132
+ memory_message_step_logger = UserMemoryMessageLogger(
133
+ message_step_logger,
129
134
  logger=logger,
130
135
  )
131
-
132
- if user_memory_state is not None:
133
- user_memory_text = user_memory_state.text
136
+ await memory_message_step_logger.log_loading_start()
137
+ user_memory_state = None
138
+ load_succeeded = False
139
+ try:
140
+ user_memory_state = await load_user_memory(
141
+ event=event,
142
+ config=user_memory_config,
143
+ language_model=memory_language_model,
144
+ logger=logger,
145
+ )
146
+ load_succeeded = True
147
+ except Exception as exc:
148
+ logger.warning(
149
+ "[user-memory] load raised - running without memory: [%s] %s",
150
+ type(exc).__name__,
151
+ exc,
152
+ )
153
+ finally:
154
+ # Always close the RUNNING step — otherwise the chat Steps UI stays stuck
155
+ # on "Loading context memory" for that turn when load raises.
156
+ if not load_succeeded:
157
+ await memory_message_step_logger.log_loading_failed()
158
+
159
+ if load_succeeded and user_memory_state is not None:
160
+ await memory_message_step_logger.log_loading_complete(with_settings_entry=True)
161
+ # The postprocessor keeps the full file (it needs the frontmatter to
162
+ # carry turn_count forward); the prompt only gets the Markdown body.
163
+ user_memory_text = profile_body(user_memory_state.text)
134
164
  postprocessor_manager.add_postprocessor(
135
165
  UserMemoryPostprocessor(
136
166
  config=user_memory_config,
@@ -138,8 +168,11 @@ if user_memory_state is not None:
138
168
  event=event,
139
169
  state=user_memory_state,
140
170
  logger=logger,
171
+ message_step_logger=memory_message_step_logger,
141
172
  )
142
173
  )
174
+ elif load_succeeded:
175
+ await memory_message_step_logger.log_loading_complete(with_settings_entry=False)
143
176
  ```
144
177
 
145
178
  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,9 @@ 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
+ - `profile_body(...)` - strips the YAML frontmatter and returns only the Markdown body. Use it whenever the profile is shown to a model; the frontmatter is bookkeeping for consolidation.
14
+ - `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
15
  - `UserMemoryPostprocessor` - runs after the assistant response, consolidates the latest turn into the profile, and uploads the updated `memory.md`.
14
16
 
15
17
  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 +19,13 @@ The memory file is intentionally small and structured. It is rewritten as a full
17
19
  ## Lifecycle
18
20
 
19
21
  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.
22
+ 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.
23
+ 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.
24
+ 4. If memory was loaded, `profile_body(...)` of its text is passed into the agent context for the current turn — the prompt only gets the Markdown body, while the postprocessor keeps the full file because it needs the frontmatter to carry `turn_count` forward.
25
+ 5. `UserMemoryPostprocessor` runs after the assistant response.
26
+ 6. The package asks the configured language model to either return `NOOP` or a complete rewritten profile.
27
+ 7. If a rewrite runs, an **Updating your memory** Step is shown while consolidating (no settings entry yet).
28
+ 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
29
 
26
30
  ## Storage Model
27
31
 
@@ -95,7 +99,9 @@ Typical orchestration code loads memory before the agent loop and registers the
95
99
  `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
100
 
97
101
  ```python
98
- from unique_user_memory.user_memory import load_user_memory
102
+ from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
103
+ from unique_user_memory.user_memory import load_user_memory, profile_body
104
+ from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
99
105
  from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
100
106
 
101
107
  user_memory_config = config.agent.services.user_memory_config
@@ -108,15 +114,39 @@ memory_language_model = (
108
114
  else user_memory_config.language_model
109
115
  )
110
116
 
111
- user_memory_state = await load_user_memory(
112
- event=event,
113
- config=user_memory_config,
114
- language_model=memory_language_model,
117
+ message_step_logger = MessageStepLogger(chat_service)
118
+ memory_message_step_logger = UserMemoryMessageLogger(
119
+ message_step_logger,
115
120
  logger=logger,
116
121
  )
117
-
118
- if user_memory_state is not None:
119
- user_memory_text = user_memory_state.text
122
+ await memory_message_step_logger.log_loading_start()
123
+ user_memory_state = None
124
+ load_succeeded = False
125
+ try:
126
+ user_memory_state = await load_user_memory(
127
+ event=event,
128
+ config=user_memory_config,
129
+ language_model=memory_language_model,
130
+ logger=logger,
131
+ )
132
+ load_succeeded = True
133
+ except Exception as exc:
134
+ logger.warning(
135
+ "[user-memory] load raised - running without memory: [%s] %s",
136
+ type(exc).__name__,
137
+ exc,
138
+ )
139
+ finally:
140
+ # Always close the RUNNING step — otherwise the chat Steps UI stays stuck
141
+ # on "Loading context memory" for that turn when load raises.
142
+ if not load_succeeded:
143
+ await memory_message_step_logger.log_loading_failed()
144
+
145
+ if load_succeeded and user_memory_state is not None:
146
+ await memory_message_step_logger.log_loading_complete(with_settings_entry=True)
147
+ # The postprocessor keeps the full file (it needs the frontmatter to
148
+ # carry turn_count forward); the prompt only gets the Markdown body.
149
+ user_memory_text = profile_body(user_memory_state.text)
120
150
  postprocessor_manager.add_postprocessor(
121
151
  UserMemoryPostprocessor(
122
152
  config=user_memory_config,
@@ -124,8 +154,11 @@ if user_memory_state is not None:
124
154
  event=event,
125
155
  state=user_memory_state,
126
156
  logger=logger,
157
+ message_step_logger=memory_message_step_logger,
127
158
  )
128
159
  )
160
+ elif load_succeeded:
161
+ await memory_message_step_logger.log_loading_complete(with_settings_entry=False)
129
162
  ```
130
163
 
131
164
  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.dev4"
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,11 @@ 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
+ profile_body,
24
25
  should_consolidate_memory,
25
26
  upload_user_memory,
26
27
  )
28
+ from unique_user_memory.user_memory_message_log import UserMemoryMessageLogger
27
29
  from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
28
30
  from unique_user_memory.user_memory_prompts import (
29
31
  consolidation_system_prompt,
@@ -45,6 +47,37 @@ def test_memory_profile_keeps_follow_up_tasks_but_excludes_open_questions() -> N
45
47
  assert "Concrete future tasks the user intends to complete" in gate_prompt
46
48
 
47
49
 
50
+ def test_profile_body_strips_frontmatter() -> None:
51
+ # The orchestrator renders this into the system prompt, where the
52
+ # bookkeeping fields are noise for the model.
53
+ content = (
54
+ "---\n"
55
+ "user_id: 233737684428787846\n"
56
+ "strategy: codex\n"
57
+ "schema_version: 1\n"
58
+ "last_updated: 2026-07-27T10:20:00+00:00\n"
59
+ "turn_count: 83\n"
60
+ "---\n\n"
61
+ "# User Memory\n\n## Identity\n- Andreas\n"
62
+ )
63
+
64
+ body = profile_body(content)
65
+
66
+ assert body == "# User Memory\n\n## Identity\n- Andreas"
67
+ for field in ("user_id:", "strategy:", "schema_version:", "turn_count:"):
68
+ assert field not in body
69
+
70
+
71
+ def test_profile_body_leaves_frontmatterless_profile_untouched() -> None:
72
+ content = "# User Memory\n\n## Identity\n- Andreas"
73
+
74
+ assert profile_body(content) == content
75
+
76
+
77
+ def test_profile_body_returns_empty_for_frontmatter_only_profile() -> None:
78
+ assert profile_body("---\nuser_id: 42\nturn_count: 0\n---\n") == ""
79
+
80
+
48
81
  def test_enforce_token_cap_truncates_long_content() -> None:
49
82
  content = "\n\n".join(f"paragraph {index} " + "word " * 40 for index in range(50))
50
83
 
@@ -679,67 +712,113 @@ async def test_consolidate_user_memory_invokes_update_end_when_start_cancelled(
679
712
 
680
713
 
681
714
  @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."
715
+ async def test_user_memory_message_step_logger_load_emits_settings_detail_entry() -> (
716
+ None
717
+ ):
718
+ step_logger = MagicMock()
719
+ step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
720
+ message_logger = UserMemoryMessageLogger(step_logger)
721
+
722
+ await message_logger.log_loading_start()
723
+ await message_logger.log_loading_complete(with_settings_entry=True)
724
+
725
+ assert step_logger.create_or_update_message_log_async.await_count == 2
726
+ start_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
727
+ 0
728
+ ].kwargs
729
+ assert start_kwargs["header"] == "Loading context memory"
730
+ assert start_kwargs["details"] is None
731
+ assert start_kwargs["references"] == []
732
+ complete_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
733
+ 1
734
+ ].kwargs
735
+ assert complete_kwargs["status"] == MessageLogStatus.COMPLETED
736
+ assert complete_kwargs["references"] == []
737
+ entry = complete_kwargs["details"].data[0]
738
+ assert entry.type == "UserMemory"
739
+ assert entry.text == "Context memory"
687
740
 
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
741
 
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
- )
742
+ @pytest.mark.asyncio
743
+ async def test_user_memory_message_step_logger_load_skips_entry_when_disabled() -> None:
744
+ step_logger = MagicMock()
745
+ step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
746
+ message_logger = UserMemoryMessageLogger(step_logger)
723
747
 
724
- await postprocessor.run(loop_response)
748
+ await message_logger.log_loading_complete(with_settings_entry=False)
725
749
 
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
750
+ complete_kwargs = step_logger.create_or_update_message_log_async.await_args.kwargs
751
+ assert complete_kwargs["details"] is None
752
+ assert complete_kwargs["references"] == []
732
753
 
733
754
 
734
755
  @pytest.mark.asyncio
735
- async def test_user_memory_postprocessor_skips_notice_when_disabled(
756
+ async def test_user_memory_message_step_logger_load_failed_marks_failed_status() -> (
757
+ None
758
+ ):
759
+ """
760
+ Purpose: log_loading_failed updates the loading Step to FAILED with no
761
+ settings entry.
762
+ Why this matters: Callers must close a RUNNING loading Step when load
763
+ raises so the chat UI does not stay stuck.
764
+ Setup summary: Start then fail the loading step; assert FAILED status.
765
+ """
766
+ step_logger = MagicMock()
767
+ step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
768
+ message_logger = UserMemoryMessageLogger(step_logger)
769
+
770
+ await message_logger.log_loading_start()
771
+ await message_logger.log_loading_failed()
772
+
773
+ failed_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
774
+ -1
775
+ ].kwargs
776
+ assert failed_kwargs["header"] == "Loading context memory"
777
+ assert failed_kwargs["status"] == MessageLogStatus.FAILED
778
+ assert failed_kwargs["details"] is None
779
+ assert failed_kwargs["references"] == []
780
+
781
+
782
+ @pytest.mark.asyncio
783
+ async def test_user_memory_message_step_logger_update_attaches_review_entry_only_when_requested() -> (
784
+ None
785
+ ):
786
+ step_logger = MagicMock()
787
+ step_logger.create_or_update_message_log_async = AsyncMock(return_value=MagicMock())
788
+ message_logger = UserMemoryMessageLogger(step_logger)
789
+
790
+ await message_logger.log_updating_start()
791
+ await message_logger.log_updating_complete(with_settings_entry=False)
792
+ await message_logger.log_updating_complete(with_settings_entry=True)
793
+
794
+ assert step_logger.create_or_update_message_log_async.await_count == 3
795
+ start_kwargs = step_logger.create_or_update_message_log_async.await_args_list[
796
+ 0
797
+ ].kwargs
798
+ assert start_kwargs["header"] == "Updating your memory"
799
+ assert start_kwargs["details"] is None
800
+ complete_without_entry = (
801
+ step_logger.create_or_update_message_log_async.await_args_list[1].kwargs
802
+ )
803
+ assert complete_without_entry["details"] is None
804
+ complete_with_entry = (
805
+ step_logger.create_or_update_message_log_async.await_args_list[2].kwargs
806
+ )
807
+ assert complete_with_entry["references"] == []
808
+ entry = complete_with_entry["details"].data[0]
809
+ assert entry.type == "UserMemory"
810
+ assert entry.text == "Review your context memory"
811
+
812
+
813
+ @pytest.mark.asyncio
814
+ async def test_user_memory_postprocessor_emits_updating_message_logs(
736
815
  monkeypatch: pytest.MonkeyPatch,
737
816
  ) -> None:
738
817
  updated_memory = "# User Memory\n\n## Identity\n- Updated"
739
818
 
740
819
  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
820
+ await on_update_start()
821
+ await on_update_end()
743
822
  return updated_memory
744
823
 
745
824
  monkeypatch.setattr(
@@ -750,32 +829,34 @@ async def test_user_memory_postprocessor_skips_notice_when_disabled(
750
829
  "unique_user_memory.user_memory_postprocessor.upload_user_memory",
751
830
  AsyncMock(return_value=True),
752
831
  )
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
- )
832
+ message_step_logger = MagicMock()
833
+ message_step_logger.log_updating_start = AsyncMock()
834
+ message_step_logger.log_updating_complete = AsyncMock()
759
835
  event = MagicMock()
760
836
  event.user_id = "user_1"
761
837
  event.company_id = "company_1"
762
838
  event.payload.user_message.text = "remember this"
763
839
  loop_response = MagicMock()
764
- loop_response.message.text = "answer"
765
- loop_response.message.id = "msg_1"
766
- loop_response.message.references = []
840
+ loop_response.message.text = "Here is your answer."
767
841
  postprocessor = UserMemoryPostprocessor(
768
- config=UserMemoryConfig(updating_notice_enabled=False),
842
+ config=UserMemoryConfig(),
769
843
  language_model=_TEST_LANGUAGE_MODEL,
770
844
  event=event,
771
- state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
845
+ state=UserMemoryState(
846
+ scope_id="scope_1",
847
+ text=empty_profile("user_1"),
848
+ ),
772
849
  logger=MagicMock(),
773
- chat_service=chat_service,
850
+ message_step_logger=message_step_logger,
774
851
  )
775
852
 
776
853
  await postprocessor.run(loop_response)
777
854
 
778
- chat_service.modify_assistant_message_async.assert_not_awaited()
855
+ message_step_logger.log_updating_start.assert_awaited_once_with()
856
+ assert [
857
+ call.kwargs
858
+ for call in message_step_logger.log_updating_complete.await_args_list
859
+ ] == [{"with_settings_entry": False}, {"with_settings_entry": True}]
779
860
 
780
861
 
781
862
  @pytest.mark.asyncio
@@ -788,14 +869,14 @@ async def test_download_user_memory_returns_empty_when_file_missing(
788
869
  search_contents,
789
870
  )
790
871
 
791
- result = await download_user_memory(
872
+ text = await download_user_memory(
792
873
  scope_id="scope_1",
793
874
  user_id="user_1",
794
875
  company_id="company_1",
795
876
  logger=MagicMock(),
796
877
  )
797
878
 
798
- assert result == ""
879
+ assert text == ""
799
880
  search_contents.assert_awaited_once_with(
800
881
  user_id="user_1",
801
882
  company_id="company_1",
@@ -822,14 +903,14 @@ async def test_download_user_memory_downloads_existing_file_to_memory(
822
903
  download_content,
823
904
  )
824
905
 
825
- result = await download_user_memory(
906
+ text = await download_user_memory(
826
907
  scope_id="scope_1",
827
908
  user_id="user_1",
828
909
  company_id="company_1",
829
910
  logger=MagicMock(),
830
911
  )
831
912
 
832
- assert result == "# User Memory\n\n## Identity\n- Test"
913
+ assert text == "# User Memory\n\n## Identity\n- Test"
833
914
  search_contents.assert_awaited_once_with(
834
915
  user_id="user_1",
835
916
  company_id="company_1",
@@ -1015,7 +1096,9 @@ async def test_ensure_user_memory_folder_returns_none_when_access_grant_fails_af
1015
1096
  async def test_upload_user_memory_writes_hidden_skip_ingestion_file(
1016
1097
  monkeypatch: pytest.MonkeyPatch,
1017
1098
  ) -> None:
1018
- upload_content = AsyncMock()
1099
+ uploaded = MagicMock()
1100
+ uploaded.id = "content_uploaded"
1101
+ upload_content = AsyncMock(return_value=uploaded)
1019
1102
  monkeypatch.setattr(
1020
1103
  "unique_user_memory.user_memory.upload_content_from_bytes_async",
1021
1104
  upload_content,
@@ -1065,14 +1148,17 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
1065
1148
  loop_response = MagicMock()
1066
1149
  loop_response.message.text = "noted"
1067
1150
  logger = MagicMock()
1068
- chat_service = MagicMock()
1151
+ message_step_logger = MagicMock(
1152
+ log_updating_start=AsyncMock(),
1153
+ log_updating_complete=AsyncMock(),
1154
+ )
1069
1155
  postprocessor = UserMemoryPostprocessor(
1070
1156
  config=UserMemoryConfig(),
1071
1157
  language_model=_TEST_LANGUAGE_MODEL,
1072
1158
  event=event,
1073
1159
  state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
1074
1160
  logger=logger,
1075
- chat_service=chat_service,
1161
+ message_step_logger=message_step_logger,
1076
1162
  )
1077
1163
 
1078
1164
  updated = await postprocessor.run(loop_response)
@@ -1085,6 +1171,9 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
1085
1171
  company_id="company_1",
1086
1172
  logger=logger,
1087
1173
  )
1174
+ message_step_logger.log_updating_complete.assert_awaited_once_with(
1175
+ with_settings_entry=True
1176
+ )
1088
1177
  logger.info.assert_any_call(
1089
1178
  "[user-memory] memory updated and uploaded successfully"
1090
1179
  )
@@ -1141,7 +1230,10 @@ async def test_user_memory_postprocessor_run_resets_invocation_stats(
1141
1230
  event=event,
1142
1231
  state=state,
1143
1232
  logger=MagicMock(),
1144
- chat_service=MagicMock(),
1233
+ message_step_logger=MagicMock(
1234
+ log_updating_start=AsyncMock(),
1235
+ log_updating_complete=AsyncMock(),
1236
+ ),
1145
1237
  )
1146
1238
 
1147
1239
  await postprocessor.run(loop_response)
@@ -1181,7 +1273,7 @@ def test_user_memory_postprocessor_take_pending_invocation_stats_drains_once() -
1181
1273
  event=event,
1182
1274
  state=state,
1183
1275
  logger=MagicMock(),
1184
- chat_service=MagicMock(),
1276
+ message_step_logger=MagicMock(),
1185
1277
  )
1186
1278
 
1187
1279
  taken = postprocessor.take_pending_invocation_stats()
@@ -1210,14 +1302,16 @@ async def test_user_memory_postprocessor_does_not_log_success_when_upload_fails(
1210
1302
  loop_response = MagicMock()
1211
1303
  loop_response.message.text = "noted"
1212
1304
  logger = MagicMock()
1213
- chat_service = MagicMock()
1214
1305
  postprocessor = UserMemoryPostprocessor(
1215
1306
  config=UserMemoryConfig(),
1216
1307
  language_model=_TEST_LANGUAGE_MODEL,
1217
1308
  event=event,
1218
1309
  state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
1219
1310
  logger=logger,
1220
- chat_service=chat_service,
1311
+ message_step_logger=MagicMock(
1312
+ log_updating_start=AsyncMock(),
1313
+ log_updating_complete=AsyncMock(),
1314
+ ),
1221
1315
  )
1222
1316
 
1223
1317
  updated = await postprocessor.run(loop_response)
@@ -63,7 +63,14 @@ _FRONTMATTER_RE = re.compile(r"^---\n.*?\n---\n", re.DOTALL)
63
63
  _TURN_COUNT_RE = re.compile(r"^turn_count:\s*(\d+)\s*$", re.MULTILINE)
64
64
 
65
65
 
66
- def _profile_body(content: str) -> str:
66
+ def profile_body(content: str) -> str:
67
+ """Return the profile without its YAML frontmatter.
68
+
69
+ The frontmatter is bookkeeping for the consolidation pass (turn count,
70
+ schema version, timestamps); consumers that show the profile to a model --
71
+ consolidation prompts and the orchestrator system prompt -- only want the
72
+ Markdown body.
73
+ """
67
74
  return _FRONTMATTER_RE.sub("", content, count=1).strip()
68
75
 
69
76
 
@@ -195,7 +202,7 @@ async def condense_user_memory(
195
202
  condensed profile, or ``None`` when the call fails or the output does
196
203
  not look like a profile (the caller then falls back to a hard cut).
197
204
  """
198
- body = _profile_body(content)
205
+ body = profile_body(content)
199
206
  current_tokens = count_tokens(content=body, language_model=language_model)
200
207
  target_tokens = max(1, int(max_tokens * _CONDENSE_TARGET_RATIO))
201
208
 
@@ -265,7 +272,7 @@ async def condense_user_memory(
265
272
  )
266
273
  return None
267
274
 
268
- candidate = _profile_body(_strip_code_fences(raw))
275
+ candidate = profile_body(_strip_code_fences(raw))
269
276
  if not _is_well_formed_profile(candidate):
270
277
  logger.warning(
271
278
  "[user-memory] condense output did not look like a profile (%d chars)",
@@ -826,7 +833,7 @@ async def _rewrite_user_memory(
826
833
  ),
827
834
  LanguageModelUserMessage(
828
835
  content=consolidation_user_prompt(
829
- existing_memory=_profile_body(safe_current),
836
+ existing_memory=profile_body(safe_current),
830
837
  user_message=_sanitize_for_xml_context(user_message or ""),
831
838
  assistant_message=_sanitize_for_xml_context(
832
839
  assistant_message or ""
@@ -883,7 +890,7 @@ async def _rewrite_user_memory(
883
890
  logger.info("[user-memory] consolidation NOOP - keeping existing memory")
884
891
  return safe_current
885
892
 
886
- candidate_body = _profile_body(_strip_code_fences(raw))
893
+ candidate_body = profile_body(_strip_code_fences(raw))
887
894
  if not _is_well_formed_profile(candidate_body):
888
895
  logger.warning(
889
896
  "[user-memory] LLM output did not look like a profile (%d chars)",
@@ -891,7 +898,7 @@ async def _rewrite_user_memory(
891
898
  )
892
899
  return safe_current
893
900
 
894
- if safe_current and candidate_body == _profile_body(safe_current):
901
+ if safe_current and candidate_body == profile_body(safe_current):
895
902
  logger.debug("[user-memory] memory body unchanged - skipping update")
896
903
  return safe_current
897
904
 
@@ -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: