Context fields may expose sensitive information if used in filtering

edwinyyyu · edwinyyyu · commit 7582233c2b33 · 2026-05-04T09:59:00.000-07:00
Signed-off-by: Edwin Yu &lt;edwinyyyu@gmail.com&gt;
diff --git a/packages/server/server_tests/memmachine_server/episodic_memory/event_memory/segment_store/test_sqlalchemy_segment_store.py b/packages/server/server_tests/memmachine_server/episodic_memory/event_memory/segment_store/test_sqlalchemy_segment_store.py
@@ -155,6 +155,12 @@ async def test_add_segments_with_message_context(
     result = await partition.get_segment_contexts([seg.uuid])
     assert result[seg.uuid][0].context == ctx
 
+    async with partition._create_session() as session:
+        row = (
+            await session.execute(select(SegmentRow).where(SegmentRow.uuid == seg.uuid))
+        ).scalar_one()
+    assert row.context == {"type": "message", "source": "User"}
+
 
 @pytest.mark.asyncio
 async def test_add_segments_with_citation_context(
@@ -386,7 +392,7 @@ async def test_contexts_property_filter(
 async def test_contexts_filter_by_context_source(
     partition: SQLAlchemySegmentStorePartition,
 ) -> None:
-    """Filter using ``context.source`` resolves via the type-tagged JSON path."""
+    """Filter using ``context.source`` is not supported."""
     ep = uuid4()
     s0 = _seg(
         event_uuid=ep,
@@ -409,22 +415,20 @@ async def test_contexts_filter_by_context_source(
     await partition.add_segments(_links(s0, s1, s2))
 
     filt = Comparison(field="context.source", op="=", value="Alice")
-    result = await partition.get_segment_contexts(
-        [s0.uuid],
-        max_backward_segments=5,
-        max_forward_segments=5,
-        property_filter=filt,
-    )
-    ctx = result[s0.uuid]
-    # s1 excluded (source=Bob); s0 seed, s2 forward.
-    assert [s.uuid for s in ctx] == [s0.uuid, s2.uuid]
+    with pytest.raises(ValueError, match="Unknown filter field"):
+        await partition.get_segment_contexts(
+            [s0.uuid],
+            max_backward_segments=5,
+            max_forward_segments=5,
+            property_filter=filt,
+        )
 
 
 @pytest.mark.asyncio
 async def test_contexts_filter_by_context_type(
     partition: SQLAlchemySegmentStorePartition,
 ) -> None:
-    """Filter using ``context.type`` discriminates between Context subtypes."""
+    """Filter using ``context.type`` is not supported."""
     ep = uuid4()
     s0 = _seg(
         event_uuid=ep,
@@ -447,15 +451,13 @@ async def test_contexts_filter_by_context_type(
     await partition.add_segments(_links(s0, s1, s2))
 
     filt = Comparison(field="context.type", op="=", value="message")
-    result = await partition.get_segment_contexts(
-        [s0.uuid],
-        max_backward_segments=5,
-        max_forward_segments=5,
-        property_filter=filt,
-    )
-    ctx = result[s0.uuid]
-    # s1 excluded (type=citation); s0 seed, s2 forward.
-    assert [s.uuid for s in ctx] == [s0.uuid, s2.uuid]
+    with pytest.raises(ValueError, match="Unknown filter field"):
+        await partition.get_segment_contexts(
+            [s0.uuid],
+            max_backward_segments=5,
+            max_forward_segments=5,
+            property_filter=filt,
+        )
 
 
 @pytest.mark.asyncio
diff --git a/packages/server/src/memmachine_server/episodic_memory/event_memory/segment_store/sqlalchemy_segment_store.py b/packages/server/src/memmachine_server/episodic_memory/event_memory/segment_store/sqlalchemy_segment_store.py
@@ -233,7 +233,7 @@ async def _insert_segments(
                 "timestamp": ensure_tz_aware(segment.timestamp),
                 "timestamp_timezone_offset": utc_offset_seconds(segment.timestamp),
                 "context": (
-                    encode_properties(segment.context.model_dump(exclude_none=True))
+                    segment.context.model_dump(mode="json")
                     if segment.context is not None
                     else None
                 ),
@@ -650,9 +650,6 @@ def _resolve_segment_field(
         """Map a filter field name to a segment column and encoding."""
         if field == "timestamp":
             return SegmentRow.timestamp.expression, "column"
-        if field.startswith("context."):
-            key = field.removeprefix("context.")
-            return SegmentRow.context[key], "properties_json"
         internal_name, is_user_metadata = normalize_filter_field(field)
         if is_user_metadata:
             key = demangle_user_metadata_key(internal_name)
@@ -665,7 +662,7 @@ def _segment_from_segment_row(
     ) -> Segment:
         """Convert a SegmentRow into a Segment."""
         context = (
-            _ContextAdapter.validate_python(decode_properties(row.context))
+            _ContextAdapter.validate_python(row.context)
             if row.context is not None
             else None
         )
