open-metadata · TeddyCr · Jul 28, 2026 · Jul 28, 2026 · Jul 28, 2026 · Jul 28, 2026
diff --git a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContextMemoryIT.java b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContextMemoryIT.java
@@ -10,7 +10,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import org.awaitility.Awaitility;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.parallel.Execution;
@@ -28,6 +30,7 @@
 import org.openmetadata.schema.entity.teams.User;
 import org.openmetadata.schema.type.EntityHistory;
 import org.openmetadata.schema.type.EntityReference;
+import org.openmetadata.sdk.client.OpenMetadataClient;
 import org.openmetadata.sdk.fluent.Users;
 import org.openmetadata.sdk.models.ListParams;
 import org.openmetadata.sdk.models.ListResponse;
@@ -59,18 +62,16 @@ public ContextMemoryIT() {
   // ABSTRACT METHOD IMPLEMENTATIONS (Required by BaseEntityIT)
   // ===================================================================
 
-  // Default to ENTITY (org-wide) visibility so the generic search-index tests have a searchable
-  // memory: only ENTITY memories are indexed — PRIVATE/SHARED are excluded at index time (see
-  // ContextMemoryRepository.isSearchIndexable / getReindexFilter). The exclusion is covered by
-  // privateAndSharedMemoriesAreExcludedFromSearchIndex below.
+  // No shareConfig, so these fall back to the default (PRIVATE) visibility. The generic
+  // search-index tests in BaseEntityIT therefore double as the regression guard that a restricted
+  // memory still reaches the index — visibility is enforced at query time, not at index time.
   @Override
   protected CreateContextMemory createMinimalRequest(TestNamespace ns) {
     return new CreateContextMemory()
         .withName(ns.prefix("context-memory"))
         .withDescription("Test context memory")
         .withQuestion("How do I find certified tables?")
-        .withAnswer("Filter the Explore page by the Certification tag.")
-        .withShareConfig(new MemoryShareConfig().withVisibility(MemoryVisibility.ENTITY));
+        .withAnswer("Filter the Explore page by the Certification tag.");
   }
 
   @Override
@@ -79,8 +80,7 @@ protected CreateContextMemory createRequest(String name, TestNamespace ns) {
         .withName(name)
         .withDescription("Test context memory")
         .withQuestion("What is the data quality SLA?")
-        .withAnswer("Critical tables must pass tests every 24 hours.")
-        .withShareConfig(new MemoryShareConfig().withVisibility(MemoryVisibility.ENTITY));
+        .withAnswer("Critical tables must pass tests every 24 hours.");
   }
 
   @Override
@@ -422,76 +422,105 @@ void post_contextMemoryWithShareConfigVisibility_200_OK(TestNamespace ns) {
     assertEquals(MemoryVisibility.PRIVATE, memory.getShareConfig().getVisibility());
   }
 
+  /**
+   * The ContextCenter serves its listing from search whenever it passes a query, filter, sort or
+   * offset — which it always does. Restricted memories must therefore reach the search index and be
+   * filtered per subject at query time: an owner and a shared principal see their own restricted
+   * memories, everyone else sees only the org-wide ones.
+   */
   @Test
-  void privateAndSharedMemoriesAreExcludedFromSearchIndex(TestNamespace ns) throws Exception {
-    ContextMemory privateMemory =
-        createEntity(memoryWithVisibility(ns, "private-mem", MemoryVisibility.PRIVATE));
-    ContextMemory sharedMemory =
-        createEntity(memoryWithVisibility(ns, "shared-mem", MemoryVisibility.SHARED));
-    // Create the ENTITY memory last. Once it is searchable, keep checking the earlier PRIVATE and
-    // SHARED memories across a stability window so delayed or out-of-order indexing cannot make the
-    // test pass before a leaked document appears.
-    ContextMemory entityMemory =
-        createEntity(memoryWithVisibility(ns, "entity-mem", MemoryVisibility.ENTITY));
-
-    Awaitility.await("ENTITY-visibility memory indexed")
-        .pollDelay(Duration.ofMillis(500))
-        .pollInterval(Duration.ofSeconds(2))
-        .atMost(Duration.ofSeconds(180))
-        .ignoreExceptions()
-        .untilAsserted(
-            () ->
-                assertTrue(
-                    searchForEntity(entityMemory.getId().toString())
-                        .contains(entityMemory.getId().toString()),
-                    "ENTITY-visibility memory must be searchable"));
+  void restrictedMemoriesAreSearchableByOwnerAndSharedPrincipalOnly(TestNamespace ns) {
+    ContextMemory ownedByUser1 =
+        createEntity(
+            memoryWithVisibility(ns, "owned-private", MemoryVisibility.PRIVATE)
+                .withOwners(List.of(testUser1Ref())));
+    ContextMemory sharedWithUser1 =
+        createEntity(
+            memoryWithVisibility(ns, "shared-with-user1", MemoryVisibility.SHARED)
+                .withShareConfig(
+                    new MemoryShareConfig()
+                        .withVisibility(MemoryVisibility.SHARED)
+                        .withSharedWith(
+                            List.of(new MemorySharedPrincipal().withPrincipal(testUser1Ref())))));
+    ContextMemory ownedByAdmin =
+        createEntity(memoryWithVisibility(ns, "admin-private", MemoryVisibility.PRIVATE));
+    ContextMemory orgWide =
+        createEntity(memoryWithVisibility(ns, "org-wide", MemoryVisibility.ENTITY));
+
+    // Admin bypasses the visibility filter, so an admin-visible listing containing all four is the
+    // barrier proving every document — restricted ones included — reached the index.
+    awaitSearchBackedListContains(
+        SdkClients.adminClient(),
+        List.of(ownedByUser1, sharedWithUser1, ownedByAdmin, orgWide),
+        "every memory must be indexed regardless of visibility");
+
+    Set<String> visibleToUser1 = searchBackedListIds(SdkClients.user1Client());
 
-    Awaitility.await("PRIVATE and SHARED memories remain excluded")
-        .pollDelay(Duration.ofMillis(500))
-        .pollInterval(Duration.ofSeconds(2))
-        .during(Duration.ofSeconds(10))
-        .atMost(Duration.ofSeconds(30))
-        .ignoreExceptions()
-        .untilAsserted(
-            () -> {
-              assertFalse(
-                  searchForEntity(privateMemory.getId().toString())
-                      .contains(privateMemory.getId().toString()),
-                  "PRIVATE memory must not be in the search index");
-              assertFalse(
-                  searchForEntity(sharedMemory.getId().toString())
-                      .contains(sharedMemory.getId().toString()),
-                  "SHARED memory must not be in the search index");
-            });
+    assertTrue(
+        visibleToUser1.contains(ownedByUser1.getId().toString()),
+        "an owner must find their own PRIVATE memory through the search-backed listing");
+    assertTrue(
+        visibleToUser1.contains(sharedWithUser1.getId().toString()),
+        "a shared principal must find the SHARED memory through the search-backed listing");
+    assertTrue(
+        visibleToUser1.contains(orgWide.getId().toString()),
+        "org-wide memories stay visible to everyone");
+    assertFalse(
+        visibleToUser1.contains(ownedByAdmin.getId().toString()),
+        "another user's PRIVATE memory must never surface");
   }
 
   @Test
-  void memoryVisibilityFlipAddsAndRemovesFromSearchIndex(TestNamespace ns) throws Exception {
+  void memoryStaysInSearchIndexAcrossVisibilityFlips(TestNamespace ns) throws Exception {
     ContextMemory memory = createEntity(memoryWithVisibility(ns, "flip", MemoryVisibility.ENTITY));
     String id = memory.getId().toString();
-    awaitSearchPresence(id, true, "ENTITY-visibility memory must be searchable");
+    awaitSearchPresence(id, "ENTITY-visibility memory must be searchable");
 
-    // Flip ENTITY -> PRIVATE: the live update must delete the doc from the index.
     ContextMemory toPrivate = getEntity(id);
     toPrivate.setShareConfig(new MemoryShareConfig().withVisibility(MemoryVisibility.PRIVATE));
     patchEntity(id, toPrivate);
-    awaitSearchPresence(
-        id, false, "Memory flipped to PRIVATE must be removed from the search index");
+    awaitSearchPresence(id, "a memory flipped to PRIVATE stays indexed for its owner");
 
-    // Flip PRIVATE -> ENTITY: the live update must re-index the doc.
     ContextMemory toEntity = getEntity(id);
     toEntity.setShareConfig(new MemoryShareConfig().withVisibility(MemoryVisibility.ENTITY));
     patchEntity(id, toEntity);
-    awaitSearchPresence(id, true, "Memory flipped back to ENTITY must be searchable again");
+    awaitSearchPresence(id, "a memory flipped back to ENTITY stays searchable");
+  }
+
+  private void awaitSearchPresence(String id, String message) {
+    Awaitility.await(message)
+        .pollDelay(Duration.ofMillis(500))
+        .pollInterval(Duration.ofSeconds(2))
+        .atMost(Duration.ofSeconds(180))
+        .ignoreExceptions()
+        .untilAsserted(() -> assertTrue(searchForEntity(id).contains(id), message));
   }
 
-  private void awaitSearchPresence(String id, boolean present, String message) {
+  private void awaitSearchBackedListContains(
+      OpenMetadataClient client, List<ContextMemory> memories, String message) {
     Awaitility.await(message)
         .pollDelay(Duration.ofMillis(500))
         .pollInterval(Duration.ofSeconds(2))
         .atMost(Duration.ofSeconds(180))
         .ignoreExceptions()
-        .untilAsserted(() -> assertEquals(present, searchForEntity(id).contains(id), message));
+        .untilAsserted(
+            () -> {
+              Set<String> listed = searchBackedListIds(client);
+              for (ContextMemory memory : memories) {
+                assertTrue(listed.contains(memory.getId().toString()), message);
+              }
+            });
+  }
+
+  /** Lists memories through the search-backed branch of the endpoint (offset forces it). */
+  private Set<String> searchBackedListIds(OpenMetadataClient client) {
+    ListParams params = new ListParams();
+    params.setLimit(1000);
+    params.setOffset(0);
+    return new ContextMemoryService(client.getHttpClient())
+        .list(params).getData().stream()
+            .map(memory -> memory.getId().toString())
+            .collect(Collectors.toSet());
   }
 
   private CreateContextMemory memoryWithVisibility(
@@ -500,7 +529,7 @@ private CreateContextMemory memoryWithVisibility(
         .withName(ns.prefix(name))
         .withDescription("Visibility indexing test")
         .withQuestion("Is this memory searchable?")
-        .withAnswer("Only when visibility is Entity.")
+        .withAnswer("Visibility is enforced at query time, not at index time.")
         .withShareConfig(new MemoryShareConfig().withVisibility(visibility));
   }
 

@@ -39,6 +39,7 @@
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.UUID;
+import java.util.function.BiPredicate;
 import java.util.stream.Collectors;
 import lombok.Getter;
 import lombok.NonNull;
@@ -741,26 +742,38 @@ public static boolean hasEntityRepository(@NonNull String entityType) {
 
   /**
    * Whether an entity instance should be written to the search index, per its repository's {@link
-   * EntityRepository#isSearchIndexable} policy (e.g. {@code ContextMemoryRepository} excludes
-   * non-org-wide memories). Defaults to {@code true} for entity types with no regular entity
-   * repository (index-only / time-series sub-entities), so the live index paths never throw on an
-   * indexable but repository-less type. Returns {@code false} when the entity or its reference is
-   * missing.
+   * EntityRepository#isSearchIndexable} policy. Defaults to {@code true} for entity types with no
+   * regular entity repository (index-only / time-series sub-entities), so the live index paths
+   * never throw on an indexable but repository-less type. Returns {@code false} when the entity or
+   * its reference is missing.
    */
   public static boolean isSearchIndexable(EntityInterface entity) {
-    if (entity == null) {
-      return false;
-    }
-    EntityReference entityReference = entity.getEntityReference();
-    if (entityReference == null) {
-      return false;
-    }
-    String entityType = entityReference.getType();
-    if (entityType == null) {
-      return false;
+    return repositoryPolicyAllows(entity, EntityRepository::isSearchIndexable);
+  }
+
+  /**
+   * Whether an entity instance should be embedded into the vector/semantic index, per its
+   * repository's {@link EntityRepository#isVectorEmbeddable} policy (e.g. {@code
+   * ContextMemoryRepository} keeps non-org-wide memories out because the vector query path carries
+   * no per-document visibility filter). Defaults and null-handling match {@link
+   * #isSearchIndexable}.
+   */
+  public static boolean isVectorEmbeddable(EntityInterface entity) {
+    return repositoryPolicyAllows(entity, EntityRepository::isVectorEmbeddable);
+  }
+
+  private static boolean repositoryPolicyAllows(
+      EntityInterface entity,
+      BiPredicate<EntityRepository<? extends EntityInterface>, EntityInterface> policy) {
+    boolean allowed = false;
+    EntityReference entityReference = entity == null ? null : entity.getEntityReference();
+    String entityType = entityReference == null ? null : entityReference.getType();
+    if (entityType != null) {
+      EntityRepository<? extends EntityInterface> repository =
+          ENTITY_REPOSITORY_MAP.get(entityType);
+      allowed = repository == null || policy.test(repository, entity);
     }
-    EntityRepository<? extends EntityInterface> repository = ENTITY_REPOSITORY_MAP.get(entityType);
-    return repository == null || repository.isSearchIndexable(entity);
+    return allowed;
   }
 
   /**

@@ -1106,6 +1106,12 @@ private String enrichWithEmbedding(
       Map<String, JsonNode> existingEmbeddingsById,
       StageStatsTracker tracker) {
     try {
+      // Per-instance gate, mirroring VectorEmbeddingHandler: the entity-type check above cannot
+      // see that an individual ContextMemory is Private/Shared, and the vector query path carries
+      // no per-document visibility filter.
+      if (!Entity.isVectorEmbeddable(entity)) {
+        return json;
+      }
       ElasticSearchVectorService vectorService = ElasticSearchVectorService.getInstance();
       if (vectorService == null) {
         return json;

@@ -1135,6 +1135,12 @@ private String enrichWithEmbedding(
       StageStatsTracker tracker,
       String stagedChunkTarget) {
     try {
+      // Per-instance gate, mirroring VectorEmbeddingHandler: the entity-type check above cannot
+      // see that an individual ContextMemory is Private/Shared, and the vector query path carries
+      // no per-document visibility filter.
+      if (!Entity.isVectorEmbeddable(entity)) {
+        return json;
+      }
       OpenSearchVectorService vectorService = OpenSearchVectorService.getInstance();
       if (vectorService == null) {
         return json;

@@ -42,6 +42,14 @@
 import org.openmetadata.service.util.EntityUtil.RelationIncludes;
 import org.openmetadata.service.util.FullyQualifiedName;
 
+/**
+ * Every memory is written to the search index regardless of its {@code shareConfig.visibility}, and
+ * privacy is enforced at query time by {@link
+ * org.openmetadata.service.search.security.ContextMemorySearchVisibility}. Indexing only org-wide
+ * memories would hide a user's own PRIVATE memories and the SHARED ones they are a principal of
+ * from {@code GET /contextCenter/memories}, which serves the ContextCenter listing from search
+ * whenever it is given a query, filter, sort or offset.
+ */
 @Slf4j
 @Repository(name = "ContextMemoryRepository")
 public class ContextMemoryRepository extends EntityRepository<ContextMemory> {
@@ -69,28 +77,20 @@ public ContextMemoryRepository() {
   }
 
   /**
-   * Only org-wide ({@link MemoryVisibility#ENTITY}) memories are searchable; PRIVATE and SHARED
-   * memories (and the {@code null}/unset default, which is PRIVATE) are kept out of the search
-   * index. This live-write check and {@link #getReindexFilter()} (the bulk-reindex equivalent
-   * expressed as a DB query) are the primary index-time enforcement points. A query-time
-   * defense-in-depth filter still protects restricted documents written before this policy was
-   * introduced. Owners and shared principals read restricted memories through the REST {@code
-   * /contextCenter/memories} endpoints (see {@link
-   * org.openmetadata.service.resources.context.ContextMemoryVisibility}).
+   * Only org-wide ({@link MemoryVisibility#ENTITY}) memories are embedded. The vector query path
+   * has no per-document visibility filter to apply — chunk documents carry neither {@code
+   * visibility} nor {@code sharedWithIds} — so a restricted memory reachable through semantic
+   * search could be returned to any caller. Keyword search does not need this restriction because
+   * {@link org.openmetadata.service.search.security.ContextMemorySearchVisibility} filters it per
+   * subject.
    */
   @Override
-  public boolean isSearchIndexable(EntityInterface entity) {
-    boolean searchable = false;
+  public boolean isVectorEmbeddable(EntityInterface entity) {
+    boolean embeddable = false;
     if (entity instanceof ContextMemory memory && memory.getShareConfig() != null) {
-      searchable = memory.getShareConfig().getVisibility() == MemoryVisibility.ENTITY;
+      embeddable = memory.getShareConfig().getVisibility() == MemoryVisibility.ENTITY;
     }
-    return searchable;
-  }
-
-  @Override
-  public ListFilter getReindexFilter() {
-    return new ListFilter(Include.ALL)
-        .addQueryParam(ListFilter.MEMORY_SEARCH_VISIBILITY_PARAM, MemoryVisibility.ENTITY.value());
+    return embeddable;
   }
 
   @Override

@@ -541,7 +541,7 @@
  @Getter protected final Set<String> allowedFields;
  public final boolean supportsSoftDelete;
  @Getter protected final boolean supportsTags;
  @Getter protected final boolean supportsOwners;
  @Getter protected final boolean supportsStyle;
  @Getter protected final boolean supportsLifeCycle;
  @Getter protected final boolean supportsCertification;
@@ -2360,19 +2360,27 @@
 
   /**
    * Whether a single entity instance should be written to the search index on the live
-   * create/update path. Defaults to true; override to keep specific instances out of the index
-   * (e.g. {@link ContextMemoryRepository} excludes non-org-wide memories). The bulk reindex applies
-   * the same rule at the DB-query level via {@link #getReindexFilter()}.
+   * create/update path. Defaults to true; override to keep specific instances out of the index.
+   * The bulk reindex applies the same rule at the DB-query level via {@link #getReindexFilter()}.
    */
   public boolean isSearchIndexable(EntityInterface entity) {
     return true;
   }
 
+  /**
+   * Whether a single entity instance should be embedded into the vector/semantic index. Defaults to
+   * {@link #isSearchIndexable}; override when an instance may be keyword-searchable but must not be
+   * reachable through the vector path (e.g. {@link ContextMemoryRepository} keeps non-org-wide
+   * memories out because vector chunks carry no visibility fields to filter on).
+   */
+  public boolean isVectorEmbeddable(EntityInterface entity) {
+    return isSearchIndexable(entity);
+  }
+
   /**
    * Filter the search reindex uses to both list and count this entity's rows. Defaults to all rows;
-   * override to keep specific instances out of the search index (e.g. {@link ContextMemoryRepository}
-   * excludes non-org-wide memories). The reader and every entity-count site share this filter so the
-   * job total matches what actually gets indexed.
+   * override to keep specific instances out of the search index. The reader and every entity-count
+   * site share this filter so the job total matches what actually gets indexed.
    */
   public ListFilter getReindexFilter() {
     return new ListFilter(Include.ALL);