fix(session): use session skill schema for patch merge (#4395)

Author: sponge225

openviking/session/compressor_v3.py

@@ -53,6 +53,10 @@
 from openviking.session.memory.utils.json_parser import JsonUtils
 from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils
 from openviking.session.memory.utils.uri import generate_uri
+from openviking.session.skill.session_skill_context_provider import (
+    SESSION_SKILL_MEMORY_TYPE,
+    load_skill_extract_registry,
+)
 from openviking.session.train import (
     Case,
     ExperienceGradientContext,
@@ -875,7 +879,8 @@ async def _get_session_skill_trainer(
             gradient_estimator=_NoopGradientEstimator(),
             policy_optimizer=PatchMergePolicyOptimizer(
                 viking_fs=viking_fs,
-                memory_type="skills",
+                memory_type=SESSION_SKILL_MEMORY_TYPE,
+                memory_registry=load_skill_extract_registry(),
             ),
             policy_updater=SkillPolicyUpdater(
                 skill_processor=self.skill_processor,

openviking/session/memory/patch_merge_context_provider.py

@@ -10,6 +10,7 @@
 
 from openviking.server.identity import RequestContext
 from openviking.session.memory.dataclass import MemoryFile, MemoryTypeSchema
+from openviking.session.memory.memory_type_registry import MemoryTypeRegistry
 from openviking.session.memory.session_extract_context_provider import (
     SessionExtractContextProvider,
 )
@@ -121,9 +122,11 @@ def __init__(
         patches: list[PatchMergePatch],
         required_file_uris: list[str] | None = None,
         output_language: str | None = None,
+        memory_registry: MemoryTypeRegistry | None = None,
     ):
         super().__init__(messages=[])
         self.memory_type = memory_type
+        self._registry = memory_registry
         self.required_file_uris = list(required_file_uris or [])
         self.patches = list(patches)
         self._output_language = output_language or _resolve_patch_output_language(self.patches)
@@ -153,7 +156,8 @@ def get_tools(self) -> list[str]:
 
     def get_memory_schemas(self, ctx: RequestContext) -> list[MemoryTypeSchema]:
         del ctx
-        schema = self._get_registry().get(self.memory_type)
+        registry = self._get_registry()
+        schema = registry.get(self.memory_type)
         if schema is None or not schema.enabled:
             raise ValueError(f"Memory schema not found or disabled: {self.memory_type}")
         return [schema]

openviking/session/train/components/policy_optimizer.py

@@ -13,6 +13,7 @@
 from openviking.session.memory.dataclass import MemoryFile, StoredLink
 from openviking.session.memory.extract_loop import ExtractLoop
 from openviking.session.memory.memory_isolation_handler import MemoryIsolationHandler
+from openviking.session.memory.memory_type_registry import MemoryTypeRegistry
 from openviking.session.memory.memory_updater import ExtractContext
 from openviking.session.memory.patch_merge_context_provider import (
     PatchMergeContextProvider,
@@ -48,6 +49,7 @@ class PatchMergePolicyOptimizer:
     viking_fs: Any = None
     vlm: Any = None
     memory_type: str = "experiences"
+    memory_registry: MemoryTypeRegistry | None = None
 
     @tracer(
         "train.policy_optimizer.patch_merge.plan",
@@ -128,6 +130,7 @@ async def _run_merge_extract_loop(
         extract_context = ExtractContext(list(context.messages or []))
         provider = PatchMergeContextProvider(
             memory_type=self.memory_type,
+            memory_registry=self.memory_registry,
             required_file_uris=_required_file_uris(gradients, policy_set),
             patches=[_gradient_to_merge_patch(gradient) for gradient in gradients],
         )
@@ -525,7 +528,7 @@ def _name_field_for_memory_type(memory_type: str) -> str:
     """Return the extra_fields key for the policy name in a given memory type."""
     if memory_type == "experiences":
         return "experience_name"
-    if memory_type == "skills":
+    if memory_type in {"skills", "session_skills"}:
         return "skill_name"
     if memory_type.endswith("s"):
         return f"{memory_type[:-1]}_name"

tests/session/memory/test_patch_merge_context_provider.py

@@ -14,6 +14,7 @@
     PatchMergeContextProvider,
     PatchMergePatch,
 )
+from openviking.session.skill.session_skill_context_provider import load_skill_extract_registry
 
 
 def _memory_file(
@@ -354,6 +355,21 @@ def test_patch_merge_context_provider_get_memory_schema_raises_for_missing_type(
         provider.get_memory_schemas(ctx=None)
 
 
+def test_patch_merge_context_provider_uses_registry_override_for_session_skills():
+    provider = PatchMergeContextProvider(
+        memory_type="session_skills",
+        memory_registry=load_skill_extract_registry(),
+        required_file_uris=[],
+        patches=[],
+    )
+
+    schemas = provider.get_memory_schemas(ctx=None)
+
+    assert len(schemas) == 1
+    assert schemas[0].memory_type == "session_skills"
+    assert schemas[0].enabled is True
+
+
 def test_patch_merge_context_provider_instruction_mentions_path_field_normalization():
     provider = PatchMergeContextProvider(
         memory_type="entities",

tests/session/train/test_train_components.py

@@ -10,6 +10,10 @@
 
 from openviking.session.memory.dataclass import MemoryFile, StoredLink
 from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils
+from openviking.session.skill.session_skill_context_provider import (
+    SESSION_SKILL_MEMORY_TYPE,
+    load_skill_extract_registry,
+)
 from openviking.session.train import (
     ContentHashPolicySnapshotter,
     DryRunPolicyUpdater,
@@ -718,7 +722,10 @@ async def run(self):
                 [],
             )
 
-    monkeypatch.setattr("openviking.session.train.components.policy_optimizer.ExtractLoop", FakeExtractLoop)
+    monkeypatch.setattr(
+        "openviking.session.train.components.policy_optimizer.ExtractLoop",
+        FakeExtractLoop,
+    )
 
     plan = await PatchMergePolicyOptimizer(viking_fs=FakeVikingFS({}), vlm=object()).plan(
         [gradient],
@@ -1074,3 +1081,79 @@ async def run(self):
     assert captured["constructed"] is True
     assert plan.metadata["patch_gradient_count"] == 1
     assert plan.items[0].after_content == "merged update"
+
+
+@pytest.mark.asyncio
+async def test_patch_merge_policy_optimizer_uses_session_skill_registry(monkeypatch):
+    from openviking.session.memory.dataclass import (
+        ResolvedOperation,
+        ResolvedOperations,
+    )
+
+    skill_uri = "viking://user/u/skills/code-review/SKILL.md"
+    policy_set = ExperienceSet(root_uri="viking://user/u/skills", policies=[])
+    gradient = PatchSemanticGradient(
+        before_file=None,
+        after_file=MemoryFile(
+            uri=skill_uri,
+            content="Use this skill to review code changes.",
+            memory_type=SESSION_SKILL_MEMORY_TYPE,
+            extra_fields={
+                "memory_type": SESSION_SKILL_MEMORY_TYPE,
+                "skill_name": "code-review",
+            },
+        ),
+        base_version=None,
+        rationale="test",
+        links=[],
+        confidence=0.9,
+        metadata={},
+    )
+    captured = {}
+
+    class FakeExtractLoop:
+        def __init__(self, **kwargs):
+            captured.update(kwargs)
+
+        async def run(self):
+            return (
+                ResolvedOperations(
+                    upsert_operations=[
+                        ResolvedOperation(
+                            old_memory_file_content=None,
+                            memory_fields={
+                                "skill_name": "code-review",
+                                "content": "Merged skill content.",
+                            },
+                            memory_type=SESSION_SKILL_MEMORY_TYPE,
+                            uris=[skill_uri],
+                        )
+                    ],
+                    delete_file_contents=[],
+                    errors=[],
+                ),
+                [],
+            )
+
+    monkeypatch.setattr(
+        "openviking.session.train.components.policy_optimizer.ExtractLoop",
+        FakeExtractLoop,
+    )
+
+    plan = await PatchMergePolicyOptimizer(
+        viking_fs=FakeVikingFS({}),
+        vlm=object(),
+        memory_type=SESSION_SKILL_MEMORY_TYPE,
+        memory_registry=load_skill_extract_registry(),
+    ).plan(
+        [gradient],
+        policy_set,
+        PatchMergePolicyOptimizerContext(request_context=fake_request_context()),
+    )
+
+    assert captured["isolation_handler"].allowed_memory_types == {SESSION_SKILL_MEMORY_TYPE}
+    assert len(plan.items) == 1
+    assert plan.items[0].memory_type == SESSION_SKILL_MEMORY_TYPE
+    assert plan.items[0].target_name == "code-review"
+    assert plan.items[0].target_uri == skill_uri
+    assert plan.items[0].after_content == "Merged skill content."