From 047521b195d7fd4c2653531f628c0f2213b971f5 Mon Sep 17 00:00:00 2001 From: ehz0ah Date: Tue, 4 Aug 2026 01:53:26 +0800 Subject: [PATCH] fix(session): persist usage across requests Store usage events durably under the session lock so a later commit can preserve relations, rich skill metadata, and active-count updates across process and request boundaries. --- openviking/server/routers/sessions.py | 6 +- openviking/session/session.py | 274 ++++++++++++++++-- .../storage/queuefs/session_commit_msg.py | 1 + tests/server/test_api_sessions.py | 105 +++++++ tests/session/test_session_usage.py | 142 +++++++++ .../session/test_session_commit_resume.py | 31 +- 6 files changed, 525 insertions(+), 34 deletions(-) diff --git a/openviking/server/routers/sessions.py b/openviking/server/routers/sessions.py index aeacaedc45..bcee1fcfc5 100644 --- a/openviking/server/routers/sessions.py +++ b/openviking/server/routers/sessions.py @@ -607,7 +607,11 @@ async def record_used( resolved_skill = dict(resolved_skill) resolved_skill["uri"] = resolve_path_variables(resolved_skill["uri"]) - session.used(contexts=resolved_contexts, skill=resolved_skill) + used_async = getattr(session, "used_async", None) + if callable(used_async): + await used_async(contexts=resolved_contexts, skill=resolved_skill) + else: + session.used(contexts=resolved_contexts, skill=resolved_skill) return Response( status="ok", result={ diff --git a/openviking/session/session.py b/openviking/session/session.py index bc6e6a77c1..bff7530357 100644 --- a/openviking/session/session.py +++ b/openviking/session/session.py @@ -76,6 +76,7 @@ _SESSION_PHASE1_LOCK_TIMEOUT_SECONDS = 30.0 _MEMORY_STEP_NAMES = ("long_term", "execution") _CUMULATIVE_CHECKPOINT_VERSION = 2 +_USAGE_EVENT_LOG_NAME = ".usage.jsonl" class _ArchiveMessagesCorruptError(ValueError): @@ -543,6 +544,47 @@ class Usage: success: bool = True timestamp: str = field(default_factory=get_current_timestamp) + def to_dict(self) -> Dict[str, Any]: + return { + "uri": self.uri, + "type": self.type, + "contribution": self.contribution, + "input": self.input, + "output": self.output, + "success": self.success, + "timestamp": self.timestamp, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "Usage": + return cls( + uri=data.get("uri", ""), + type=data.get("type", ""), + contribution=data.get("contribution", 0.0), + input=data.get("input", ""), + output=data.get("output", ""), + success=data.get("success", True), + timestamp=data.get("timestamp", get_current_timestamp()), + ) + + +@dataclass(frozen=True) +class _UsageEvent: + """One durable usage event pending consumption by Session Phase 1.""" + + event_id: str + usage: Usage + + def to_dict(self) -> Dict[str, Any]: + return {"event_id": self.event_id, **self.usage.to_dict()} + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "_UsageEvent": + event_id = data.get("event_id") + if not isinstance(event_id, str) or not event_id: + raise ValueError("usage event is missing event_id") + return cls(event_id=event_id, usage=Usage.from_dict(data)) + class Session: """Session management class - Message = role + parts.""" @@ -576,6 +618,7 @@ def __init__( self._messages: List[Message] = [] self._usage_records: List[Usage] = [] + self._pending_usage_records: List[Usage] = [] self._archive_meta_merge_lock = asyncio.Lock() self._compression: SessionCompression = SessionCompression() self._stats: SessionStats = SessionStats() @@ -650,6 +693,7 @@ async def load(self): self._meta.created_by_account_id = self.ctx.account_id if not self._meta.created_by_user_id: self._meta.created_by_user_id = self.ctx.user.user_id + await self._reload_usage_records() # WM v2: always rebuild pending_tokens from current messages so the # counter stays consistent across restarts and is also backfilled for # legacy sessions whose .meta.json predates these fields. O(n) once, @@ -743,39 +787,172 @@ def used( contexts: Optional[List[str]] = None, skill: Optional[Dict[str, Any]] = None, ) -> None: - """Record actually used contexts and skills.""" - if contexts: - for uri in contexts: - usage = Usage(uri=uri, type="context") - self._usage_records.append(usage) - self._stats.contexts_used += 1 - logger.debug(f"Tracked context usage: {uri}") - try: - from openviking.metrics.datasources.session import SessionLifecycleDataSource + """Record usage in memory for direct in-process compatibility.""" + records = self._build_usage_records(contexts=contexts, skill=skill) + self._pending_usage_records.extend(records) + self._usage_records.extend(records) + self._update_usage_stats() + self._record_usage_metrics(records) + + async def used_async( + self, + contexts: Optional[List[str]] = None, + skill: Optional[Dict[str, Any]] = None, + ) -> None: + """Durably append usage records under the Session Phase 1 lock.""" + records = self._build_usage_records(contexts=contexts, skill=skill) + if not records: + return + if not self._viking_fs: + self.used(contexts=contexts, skill=skill) + return - SessionLifecycleDataSource.record_contexts_used( - action="context", delta=len(contexts) - ) - except Exception: - pass + uri_to_path = getattr(self._viking_fs, "_uri_to_path", None) + if not callable(uri_to_path): + events = await self._read_usage_events() + events.extend(self._new_usage_events(records)) + await self._write_usage_events(events) + self._replace_usage_records( + [event.usage for event in events] + list(self._pending_usage_records) + ) + self._record_usage_metrics(records) + return + session_path = uri_to_path(self._session_uri, ctx=self.ctx) + lease = await self._viking_fs._async_agfs.pathlock_acquire_tree( + session_path, timeout_secs=_SESSION_PHASE1_LOCK_TIMEOUT_SECONDS + ) + try: + events = await self._read_usage_events() + events.extend(self._new_usage_events(records)) + await self._write_usage_events(events, lease_ref=lease) + self._replace_usage_records( + [event.usage for event in events] + list(self._pending_usage_records) + ) + finally: + await self._viking_fs._async_agfs.pathlock_release(lease) + self._record_usage_metrics(records) + + def _build_usage_records( + self, + *, + contexts: Optional[List[str]], + skill: Optional[Dict[str, Any]], + ) -> List[Usage]: + records = [Usage(uri=uri, type="context") for uri in contexts or []] if skill: - usage = Usage( - uri=skill.get("uri", ""), - type="skill", - input=skill.get("input", ""), - output=skill.get("output", ""), - success=skill.get("success", True), + records.append( + Usage( + uri=skill.get("uri", ""), + type="skill", + contribution=skill.get("contribution", 0.0), + input=skill.get("input", ""), + output=skill.get("output", ""), + success=skill.get("success", True), + timestamp=skill.get("timestamp", get_current_timestamp()), + ) ) - self._usage_records.append(usage) - self._stats.skills_used += 1 - logger.debug(f"Tracked skill usage: {skill.get('uri')}") - try: - from openviking.metrics.datasources.session import SessionLifecycleDataSource + return records - SessionLifecycleDataSource.record_contexts_used(action="skill", delta=1) - except Exception: - pass + @staticmethod + def _new_usage_events(records: List[Usage]) -> List[_UsageEvent]: + return [_UsageEvent(event_id=str(uuid4()), usage=record) for record in records] + + def _usage_event_log_uri(self) -> str: + return f"{self._session_uri}/{_USAGE_EVENT_LOG_NAME}" + + async def _read_usage_events(self) -> List[_UsageEvent]: + try: + content = await self._viking_fs.read_file( + self._usage_event_log_uri(), + ctx=self.ctx, + ) + except Exception as exc: + if _is_storage_not_found(exc): + return [] + raise + return [ + _UsageEvent.from_dict(json.loads(line)) for line in content.splitlines() if line.strip() + ] + + async def _write_usage_events( + self, + events: List[_UsageEvent], + *, + lease_ref: Optional[Any] = None, + ) -> None: + content = "".join( + json.dumps(event.to_dict(), ensure_ascii=False) + "\n" for event in events + ) + write_kwargs: Dict[str, Any] = {"ctx": self.ctx} + if lease_ref is not None: + write_kwargs["lease_ref"] = lease_ref + await self._viking_fs.write_file( + self._usage_event_log_uri(), + content, + **write_kwargs, + ) + + async def _reload_usage_records(self) -> None: + events = await self._read_usage_events() + self._pending_usage_records = [] + self._replace_usage_records([event.usage for event in events]) + + def _replace_usage_records(self, records: List[Usage]) -> None: + self._usage_records = list(records) + self._update_usage_stats() + + def _update_usage_stats(self) -> None: + self._stats.contexts_used = sum( + 1 for usage in self._usage_records if usage.type == "context" + ) + self._stats.skills_used = sum(1 for usage in self._usage_records if usage.type == "skill") + + @staticmethod + def _record_usage_metrics(records: List[Usage]) -> None: + try: + from openviking.metrics.datasources.session import SessionLifecycleDataSource + + context_count = sum(1 for usage in records if usage.type == "context") + skill_count = sum(1 for usage in records if usage.type == "skill") + if context_count: + SessionLifecycleDataSource.record_contexts_used( + action="context", delta=context_count + ) + if skill_count: + SessionLifecycleDataSource.record_contexts_used(action="skill", delta=skill_count) + except Exception: + pass + + async def _consume_usage_events( + self, + event_ids: List[str], + *, + lease_ref: Optional[Any] = None, + ) -> List[_UsageEvent]: + consumed_ids = set(event_ids) + remaining = [ + event for event in await self._read_usage_events() if event.event_id not in consumed_ids + ] + await self._write_usage_events(remaining, lease_ref=lease_ref) + self._replace_usage_records( + [event.usage for event in remaining] + list(self._pending_usage_records) + ) + return remaining + + async def _restore_usage_events( + self, + snapshot: List[_UsageEvent], + *, + lease_ref: Optional[Any] = None, + ) -> None: + current = await self._read_usage_events() + current_ids = {event.event_id for event in current} + restored = [event for event in snapshot if event.event_id not in current_ids] + current + await self._write_usage_events(restored, lease_ref=lease_ref) + self._replace_usage_records( + [event.usage for event in restored] + list(self._pending_usage_records) + ) def _tool_result_store(self) -> Optional[ToolResultStore]: if not self._viking_fs: @@ -1501,6 +1678,7 @@ async def _write_phase1_marker( min_raw_tail_steps: int, agent_evolution_enabled: bool = True, agent_memory_skip_reason: Optional[str] = None, + usage_event_ids: Optional[List[str]] = None, lease_ref: Optional[Any] = None, ) -> None: """Persist the Phase 1 intent before any destructive root rewrite.""" @@ -1517,6 +1695,7 @@ async def _write_phase1_marker( "keep_recent_turn_count": keep_recent_turn_count, "retained_message_token_budget": retained_message_token_budget, "min_raw_tail_steps": min_raw_tail_steps, + "usage_event_ids": list(usage_event_ids or []), } await self._merge_archive_meta( archive_uri, @@ -1644,6 +1823,12 @@ async def _ensure_phase1_ready(self, archive_uri: str) -> bool: ) self._meta.last_commit_at = get_current_timestamp() self._rebuild_pending_tokens() + usage_event_ids = marker.get("usage_event_ids") + if isinstance(usage_event_ids, list): + await self._consume_usage_events( + [event_id for event_id in usage_event_ids if isinstance(event_id, str)], + lease_ref=lease, + ) await self._save_meta(lease_ref=lease) await self._write_phase1_ready_marker(archive_uri, lease_ref=lease) logger.warning("Recovered interrupted Session Phase 1: %s", archive_uri) @@ -1888,12 +2073,24 @@ async def commit_async( "budget_exceeded": retention_plan.budget_exceeded if retention_plan else False, } + original_messages = list(self._messages) + usage_events_snapshot = await self._read_usage_events() + pending_usage_snapshot = list(self._pending_usage_records) + if pending_usage_snapshot: + usage_events_snapshot.extend(self._new_usage_events(pending_usage_snapshot)) + await self._write_usage_events(usage_events_snapshot, lease_ref=lease) + snapshot_object_ids = {id(record) for record in pending_usage_snapshot} + self._pending_usage_records = [ + record + for record in self._pending_usage_records + if id(record) not in snapshot_object_ids + ] + usage_snapshot = [event.usage for event in usage_events_snapshot] + self._replace_usage_records(usage_snapshot + list(self._pending_usage_records)) self._compression.compression_index += 1 archive_uri = ( f"{self._session_uri}/history/archive_{self._compression.compression_index:03d}" ) - original_messages = list(self._messages) - usage_snapshot = self._usage_records.copy() task_id = str(uuid4()) queue_msg = SessionCommitMsg( task_id=task_id, @@ -1903,6 +2100,7 @@ async def commit_async( user=self.ctx.user.to_dict(), memory_policy=effective_memory_policy, usage_uris=list(dict.fromkeys(u.uri for u in usage_snapshot if u.uri)), + usage_records=[usage.to_dict() for usage in usage_snapshot], ) phase1_stage = "phase1_persist" try: @@ -1922,6 +2120,7 @@ async def commit_async( min_raw_tail_steps=effective_min_tail, agent_evolution_enabled=agent_evolution_enabled, agent_memory_skip_reason=agent_memory_skip_reason, + usage_event_ids=[event.event_id for event in usage_events_snapshot], lease_ref=lease, ) @@ -1986,6 +2185,10 @@ async def commit_async( self._compression.compression_index, ) self._meta.last_commit_at = get_current_timestamp() + await self._consume_usage_events( + [event.event_id for event in usage_events_snapshot], + lease_ref=lease, + ) await self._save_meta(lease_ref=lease) await self._write_phase1_ready_marker(archive_uri, lease_ref=lease) except Exception as e: @@ -2006,6 +2209,7 @@ async def commit_async( archive_uri, ) self._messages = original_messages + await self._restore_usage_events(usage_events_snapshot, lease_ref=lease) self._compression.compression_index -= 1 raise finally: @@ -2161,7 +2365,11 @@ async def resume_queued_commit(self, msg: "SessionCommitMsg") -> None: task_id=msg.task_id, archive_uri=msg.archive_uri, messages=archive_messages, - usage_records=[Usage(uri=uri, type="context") for uri in msg.usage_uris], + usage_records=( + [Usage.from_dict(record) for record in msg.usage_records] + if msg.usage_records + else [Usage(uri=uri, type="context") for uri in msg.usage_uris] + ), first_message_id=archive_messages[0].id, last_message_id=archive_messages[-1].id, memory_policy=msg.memory_policy, @@ -2571,6 +2779,8 @@ async def _run_execution_memory_extraction() -> Any: # Write relations (using snapshot, not self._usage_records) if self._viking_fs: for usage in usage_records: + if not usage.uri: + continue try: await self._viking_fs.link( self._session_uri, usage.uri, ctx=self.ctx @@ -2580,7 +2790,7 @@ async def _run_execution_memory_extraction() -> Any: # Update active_count (using snapshot, not self._usage_records) if self._vikingdb_manager: - uris = [u.uri for u in usage_records if u.uri] + uris = list(dict.fromkeys(u.uri for u in usage_records if u.uri)) try: active_count_updated = ( await self._vikingdb_manager.increment_active_count(self.ctx, uris) diff --git a/openviking/storage/queuefs/session_commit_msg.py b/openviking/storage/queuefs/session_commit_msg.py index 9a0efe31e5..6af505f093 100644 --- a/openviking/storage/queuefs/session_commit_msg.py +++ b/openviking/storage/queuefs/session_commit_msg.py @@ -15,6 +15,7 @@ class SessionCommitMsg: user: Dict[str, str] memory_policy: Dict[str, Any] = field(default_factory=dict) usage_uris: List[str] = field(default_factory=list) + usage_records: List[Dict[str, Any]] = field(default_factory=list) def to_dict(self) -> Dict[str, Any]: return asdict(self) diff --git a/tests/server/test_api_sessions.py b/tests/server/test_api_sessions.py index 29f3c39ce1..dfcfcbc299 100644 --- a/tests/server/test_api_sessions.py +++ b/tests/server/test_api_sessions.py @@ -5,6 +5,7 @@ import asyncio import json +import threading from unittest.mock import patch import httpx @@ -885,6 +886,110 @@ async def test_commit_updates_archive_metadata_before_background_task(client: ht assert after_result["commit_count"] == 1 +async def test_used_then_separate_commit_updates_active_count_and_relation( + client, + service, + monkeypatch, +): + lock = threading.Lock() + mock_agfs = service.viking_fs._async_agfs._client + + def acquire_tree(_ctx, _path, _timeout_secs, _owner_lease_ref): + lock.acquire() + return {"lease_ref": "session-usage-test-lock"} + + def release_tree(_ctx, _lease): + lock.release() + + monkeypatch.setattr(mock_agfs, "pathlock_acquire_tree", acquire_tree, raising=False) + monkeypatch.setattr(mock_agfs, "pathlock_release", release_tree, raising=False) + + async def no_memories(*args, **kwargs): + del args, kwargs + return [] + + async def fake_summary(*args, **kwargs): + del args, kwargs + return "# Durable Usage\n\nTest summary." + + service.sessions._session_compressor.extract_long_term_memories = no_memories + service.sessions._session_compressor.extract_execution_memories = no_memories + monkeypatch.setattr( + "openviking.session.session.Session._generate_archive_summary_async", + fake_summary, + ) + + resource_uri = "viking://resources/http-used-durability.md" + ctx = RequestContext(user=UserIdentifier.the_default_user(), role=Role.ROOT) + vector = service.vikingdb_manager.get_embedder().embed("durable usage").dense_vector + await service.vikingdb_manager.upsert( + { + "uri": resource_uri, + "parent_uri": "viking://resources", + "is_leaf": True, + "abstract": "Durable usage test resource", + "context_type": "resource", + "category": "", + "active_count": 0, + "vector": vector, + "meta": {}, + "related_uri": [], + "account_id": "default", + "owner_space": "", + "level": 2, + }, + ctx=ctx, + ) + + create_resp = await client.post( + "/api/v1/sessions", + json={"session_id": "http-used-durability"}, + ) + assert create_resp.status_code == 200 + session_id = create_resp.json()["result"]["session_id"] + + add_resp = await client.post( + f"/api/v1/sessions/{session_id}/messages", + json=_message_request("user", content="use the resource"), + ) + assert add_resp.status_code == 200 + + used_resp = await client.post( + f"/api/v1/sessions/{session_id}/used", + json={"contexts": [resource_uri]}, + ) + assert used_resp.status_code == 200 + assert used_resp.json()["result"]["contexts_used"] == 1 + + commit_resp = await client.post(f"/api/v1/sessions/{session_id}/commit") + assert commit_resp.status_code == 200 + task = await _wait_for_task(client, commit_resp.json()["result"]["task_id"]) + + assert task["status"] == "completed" + assert task["result"]["active_count_updated"] == 1 + + records = await service.vikingdb_manager.get_context_by_uri( + uri=resource_uri, + limit=1, + ctx=ctx, + ) + assert records + assert records[0]["active_count"] == 1 + + session_uri = create_resp.json()["result"]["uri"] + relations_resp = await client.get( + "/api/v1/relations", + params={"uri": session_uri}, + ) + assert relations_resp.status_code == 200 + relation_uris = { + relation["uri"] + for relation in relations_resp.json()["result"] + if isinstance(relation, dict) and "uri" in relation + } + assert resource_uri in relation_uris + + async def test_extract_session_jsonable_regression(client: httpx.AsyncClient, service, monkeypatch): """Regression: extract endpoint should serialize internal objects.""" diff --git a/tests/session/test_session_usage.py b/tests/session/test_session_usage.py index d138b2f6ba..a5e0af7242 100644 --- a/tests/session/test_session_usage.py +++ b/tests/session/test_session_usage.py @@ -3,6 +3,11 @@ """Usage record tests""" +import asyncio + +import pytest + +from openviking import AsyncOpenViking from openviking.message import TextPart from openviking.session import Session @@ -64,3 +69,140 @@ async def test_used_empty(self, session: Session): session.used() # Should not raise error + + async def test_used_async_survives_reload_with_full_usage_fields( + self, + client: AsyncOpenViking, + ): + session = client.session(session_id="durable_usage_reload") + await session.ensure_exists() + + await session.used_async( + skill={ + "uri": "viking://user/skills/search", + "input": "query", + "output": "result", + "success": False, + "contribution": 0.75, + "timestamp": "2026-08-04T01:02:03+00:00", + } + ) + + reloaded = client.session(session_id=session.session_id) + await reloaded.load() + + assert len(reloaded.usage_records) == 1 + usage = reloaded.usage_records[0] + assert usage.uri == "viking://user/skills/search" + assert usage.type == "skill" + assert usage.contribution == 0.75 + assert usage.input == "query" + assert usage.output == "result" + assert usage.success is False + assert usage.timestamp == "2026-08-04T01:02:03+00:00" + assert reloaded.stats.skills_used == 1 + + async def test_usage_append_during_phase1_remains_for_next_commit( + self, + client: AsyncOpenViking, + monkeypatch, + ): + session = client.session(session_id="durable_usage_commit_race") + session.add_message("user", [TextPart("first turn")]) + await session.used_async( + skill={ + "uri": "viking://user/skills/first", + "contribution": 0.5, + "input": "first input", + "output": "first output", + "success": False, + "timestamp": "2026-08-04T02:03:04+00:00", + } + ) + + concurrent_session = client.session(session_id=session.session_id) + await concurrent_session.load() + + phase1_snapshot_ready = asyncio.Event() + release_phase1 = asyncio.Event() + captured_usage_records = [] + original_write_phase1_marker = session._write_phase1_marker + + async def blocking_write_phase1_marker(*args, **kwargs): + captured_usage_records.extend(kwargs["queue_message"]["usage_records"]) + phase1_snapshot_ready.set() + await release_phase1.wait() + return await original_write_phase1_marker(*args, **kwargs) + + monkeypatch.setattr(session, "_write_phase1_marker", blocking_write_phase1_marker) + + commit_task = asyncio.create_task(session.commit_async()) + await phase1_snapshot_ready.wait() + append_task = asyncio.create_task( + concurrent_session.used_async(contexts=["viking://resources/second"]) + ) + await asyncio.sleep(0) + assert not append_task.done() + + release_phase1.set() + result = await commit_task + await append_task + + assert captured_usage_records == [ + { + "uri": "viking://user/skills/first", + "type": "skill", + "contribution": 0.5, + "input": "first input", + "output": "first output", + "success": False, + "timestamp": "2026-08-04T02:03:04+00:00", + } + ] + + reloaded = client.session(session_id=session.session_id) + await reloaded.load() + assert [usage.uri for usage in reloaded.usage_records] == ["viking://resources/second"] + assert result["archived"] is True + + async def test_phase1_failure_restores_usage_sidecar( + self, + client: AsyncOpenViking, + monkeypatch, + ): + session = client.session(session_id="durable_usage_phase1_failure") + session.add_message("user", [TextPart("archive candidate")]) + await session.used_async( + skill={ + "uri": "viking://user/skills/retry", + "input": "retry input", + "output": "retry output", + "success": False, + } + ) + + class FailingQueueManager: + async def enqueue(self, _queue_name, _data): + raise RuntimeError("queue unavailable") + + monkeypatch.setattr( + "openviking.storage.queuefs.get_queue_manager", + lambda: FailingQueueManager(), + ) + + with pytest.raises(RuntimeError, match="queue unavailable"): + await session.commit_async() + + reloaded = client.session(session_id=session.session_id) + await reloaded.load() + assert [usage.to_dict() for usage in reloaded.usage_records] == [ + { + "uri": "viking://user/skills/retry", + "type": "skill", + "contribution": 0.0, + "input": "retry input", + "output": "retry output", + "success": False, + "timestamp": session.usage_records[0].timestamp, + } + ] diff --git a/tests/unit/session/test_session_commit_resume.py b/tests/unit/session/test_session_commit_resume.py index 01c7b2b8c1..4c68cf3be3 100644 --- a/tests/unit/session/test_session_commit_resume.py +++ b/tests/unit/session/test_session_commit_resume.py @@ -44,7 +44,7 @@ async def read_file(self, uri, ctx=None): raise FileNotFoundError(uri) return self.files[uri] - async def write_file(self, uri, content, ctx=None): + async def write_file(self, uri, content, ctx=None, lease_ref=None): self.files[uri] = content @@ -72,6 +72,17 @@ async def test_resume_queued_commit_continues_phase2(monkeypatch): session_uri=session_uri, archive_uri=archive_uri, user={"account_id": "default", "user_id": "default"}, + usage_records=[ + { + "uri": "viking://user/skills/search", + "type": "skill", + "contribution": 0.75, + "input": "query", + "output": "result", + "success": False, + "timestamp": "2026-08-04T01:02:03+00:00", + } + ], ) try: @@ -85,6 +96,8 @@ async def test_resume_queued_commit_continues_phase2(monkeypatch): assert [ item.id for item in session._run_memory_extraction.await_args.kwargs["messages"] ] == ["archived"] + usage_records = session._run_memory_extraction.await_args.kwargs["usage_records"] + assert [usage.to_dict() for usage in usage_records] == message.usage_records @pytest.mark.asyncio @@ -209,3 +222,19 @@ def test_session_commit_message_ignores_unknown_fields(): assert message.task_id == "task-1" assert "actor_peer_id" not in message.to_dict() + + +def test_session_commit_message_accepts_legacy_usage_uris(): + message = SessionCommitMsg.from_dict( + { + "task_id": "task-1", + "session_id": "session-1", + "session_uri": "viking://user/sessions/session-1", + "archive_uri": "viking://user/sessions/session-1/history/archive_001", + "user": {"account_id": "default", "user_id": "default"}, + "usage_uris": ["viking://resources/legacy"], + } + ) + + assert message.usage_uris == ["viking://resources/legacy"] + assert message.usage_records == []