fix(workspaces): preserve file workspace memberships on PUT and PATCH

openrag/components/indexer/indexer.py

@@ -195,9 +195,19 @@ async def update_file_metadata(
             for doc in docs:
                 doc.metadata.update(metadata)
 
+            # Snapshot workspace memberships before deletion so they can be restored.
+            workspace_ids = await vectordb.get_file_workspaces.remote(file_id, partition)
+
             await self.delete_file(file_id, partition)
             await vectordb.async_add_documents.remote(docs, user=user)
 
+            # Restore workspace memberships that existed before the delete.
+            if workspace_ids:
+                await asyncio.gather(
+                    *[vectordb.add_files_to_workspace.remote(ws_id, [file_id]) for ws_id in workspace_ids]
+                )
+                log.debug("Restored workspace memberships after metadata update.", workspace_ids=workspace_ids)
+
             log.info("Metadata updated for file.")
         except Exception as e:
             log.error("Error in update_file_metadata", error=str(e))

openrag/components/indexer/vectordb/utils.py

@@ -675,6 +675,18 @@ def list_workspace_files(self, workspace_id: str) -> list[str]:
             result = session.execute(select(WorkspaceFile.file_id).where(WorkspaceFile.workspace_id == workspace_id))
             return [r[0] for r in result.all()]
 
+    def get_file_workspaces(self, file_id: str, partition: str) -> list[str]:
+        """Return the workspace IDs that contain the given file, scoped to the given partition."""
+        with self.Session() as session:
+            ws_ids = select(Workspace.workspace_id).where(Workspace.partition_name == partition)
+            result = session.execute(
+                select(WorkspaceFile.workspace_id).where(
+                    WorkspaceFile.file_id == file_id,
+                    WorkspaceFile.workspace_id.in_(ws_ids),
+                )
+            )
+            return [r[0] for r in result.all()]
+
     def remove_file_from_all_workspaces(self, file_id: str, partition: str):
         """Remove file from all workspaces in the given partition — called during file deletion."""
         with self.Session() as session:

openrag/components/indexer/vectordb/vectordb.py

@@ -1132,6 +1132,10 @@ async def remove_file_from_workspace(self, workspace_id: str, file_id: str) -> b
     async def list_workspace_files(self, workspace_id: str) -> list[str]:
         return self.partition_file_manager.list_workspace_files(workspace_id)
 
+    async def get_file_workspaces(self, file_id: str, partition: str) -> list[str]:
+        """Return workspace IDs that contain the given file, scoped to the partition."""
+        return self.partition_file_manager.get_file_workspaces(file_id, partition)
+
 
 def _gen_chunk_order_metadata(n: int = 20) -> list[dict]:
     # Use base timestamp + index to ensure uniqueness

openrag/routers/indexer.py

@@ -265,6 +265,9 @@ async def put_file(
             detail=f"'{file_id}' not found in partition '{partition}'",
         )
 
+    # Snapshot workspace memberships before deletion so they can be restored on the new version.
+    existing_workspace_ids = await vectordb.get_file_workspaces.remote(file_id, partition)
+
     # Delete the existing file from the vector database
     await indexer.delete_file.remote(file_id, partition)
 
@@ -288,8 +291,14 @@ async def put_file(
     metadata["created_at"] = datetime.fromtimestamp(file_stat.st_ctime).isoformat()
     metadata["file_id"] = file_id
 
-    # Indexing the file
-    task = indexer.add_file.remote(path=file_path, metadata=metadata, partition=partition, user=user)
+    # Indexing the file — restore pre-existing workspace memberships on the new version.
+    task = indexer.add_file.remote(
+        path=file_path,
+        metadata=metadata,
+        partition=partition,
+        user=user,
+        workspace_ids=existing_workspace_ids or None,
+    )
     await task_state_manager.set_state.remote(task.task_id().hex(), "QUEUED")
     await task_state_manager.set_object_ref.remote(task.task_id().hex(), {"ref": task})
 

openrag/routers/workspaces.py

@@ -169,6 +169,19 @@ async def list_workspace_files(
     return {"file_ids": file_ids}
 
 
+@router.get(
+    "/partition/{partition}/files/{file_id}/workspaces",
+    dependencies=[Depends(require_partition_viewer)],
+)
+async def list_file_workspaces(partition: str, file_id: str, vectordb=Depends(get_vectordb)):
+    workspace_ids = await call_ray_actor_with_timeout(
+        vectordb.get_file_workspaces.remote(file_id, partition),
+        timeout=VECTORDB_TIMEOUT,
+        task_description=f"get_file_workspaces({file_id})",
+    )
+    return {"file_id": file_id, "workspace_ids": workspace_ids}
+
+
 @router.delete(
     "/partition/{partition}/workspaces/{workspace_id}/files/{file_id}",
     dependencies=[Depends(require_partition_editor)],

tests/api_tests/test_workspaces.py

@@ -1,9 +1,12 @@
 """API integration tests for workspace endpoints."""
 
+import io
 import uuid
 
 import pytest
 
+from .conftest import wait_for_indexing
+
 pytestmark = pytest.mark.integration
 
 
@@ -165,3 +168,115 @@ def test_delete_partition_cascades_workspaces(self, api_client):
 
         # Cleanup
         api_client.delete(f"/partition/{partition}")
+
+
+class TestFileWorkspaceMembership:
+    """Test that file workspace memberships survive file replace operations."""
+
+    def _upload_file(self, api_client, partition: str, file_id: str, content: str = "Test content"):
+        file_obj = io.BytesIO(content.encode())
+        return api_client.post(
+            f"/indexer/partition/{partition}/file/{file_id}",
+            files={"file": (f"{file_id}.txt", file_obj, "text/plain")},
+            data={"metadata": "{}"},
+        )
+
+    def _get_file_workspaces(self, api_client, partition: str, file_id: str) -> list[str]:
+        response = api_client.get(f"/partition/{partition}/files/{file_id}/workspaces")
+        assert response.status_code == 200
+        return response.json()["workspace_ids"]
+
+    def test_workspace_memberships_preserved_after_put(self, api_client, workspace_partition):
+        """After a PUT (file replace), the file's workspace memberships must be intact."""
+        file_id = f"file-{uuid.uuid4().hex[:8]}"
+        ws1 = f"ws-{uuid.uuid4().hex[:8]}"
+        ws2 = f"ws-{uuid.uuid4().hex[:8]}"
+        ws3 = f"ws-{uuid.uuid4().hex[:8]}"
+
+        # Create 3 workspaces
+        for ws in [ws1, ws2, ws3]:
+            r = api_client.post(
+                f"/partition/{workspace_partition}/workspaces",
+                json={"workspace_id": ws},
+            )
+            assert r.status_code == 201
+
+        # Upload and index the file
+        response = self._upload_file(api_client, workspace_partition, file_id)
+        assert response.status_code in [200, 201, 202]
+        wait_for_indexing(api_client, response.json())
+
+        # Add the file to ws1 and ws2 (not ws3)
+        for ws in [ws1, ws2]:
+            r = api_client.post(
+                f"/partition/{workspace_partition}/workspaces/{ws}/files",
+                json={"file_ids": [file_id]},
+            )
+            assert r.status_code == 200
+
+        # Verify initial membership
+        ws_before = self._get_file_workspaces(api_client, workspace_partition, file_id)
+        assert set(ws_before) == {ws1, ws2}
+
+        # Replace the file via PUT
+        replace_obj = io.BytesIO(b"Updated content")
+        r = api_client.put(
+            f"/indexer/partition/{workspace_partition}/file/{file_id}",
+            files={"file": (f"{file_id}.txt", replace_obj, "text/plain")},
+            data={"metadata": "{}"},
+        )
+        assert r.status_code in [200, 201, 202]
+        wait_for_indexing(api_client, r.json())
+
+        # Workspace memberships must be restored
+        ws_after = self._get_file_workspaces(api_client, workspace_partition, file_id)
+        assert set(ws_after) == {ws1, ws2}, (
+            f"Expected workspace memberships {{{ws1}, {ws2}}} after PUT, got {set(ws_after)}"
+        )
+        assert ws3 not in ws_after
+
+    def test_workspace_memberships_preserved_after_patch(self, api_client, workspace_partition):
+        """After a PATCH (metadata update), the file's workspace memberships must be intact."""
+        file_id = f"file-{uuid.uuid4().hex[:8]}"
+        ws1 = f"ws-{uuid.uuid4().hex[:8]}"
+        ws2 = f"ws-{uuid.uuid4().hex[:8]}"
+        ws3 = f"ws-{uuid.uuid4().hex[:8]}"
+
+        # Create 3 workspaces
+        for ws in [ws1, ws2, ws3]:
+            r = api_client.post(
+                f"/partition/{workspace_partition}/workspaces",
+                json={"workspace_id": ws},
+            )
+            assert r.status_code == 201
+
+        # Upload and index the file
+        response = self._upload_file(api_client, workspace_partition, file_id)
+        assert response.status_code in [200, 201, 202]
+        wait_for_indexing(api_client, response.json())
+
+        # Add the file to ws1 and ws2 (not ws3)
+        for ws in [ws1, ws2]:
+            r = api_client.post(
+                f"/partition/{workspace_partition}/workspaces/{ws}/files",
+                json={"file_ids": [file_id]},
+            )
+            assert r.status_code == 200
+
+        # Verify initial membership
+        ws_before = self._get_file_workspaces(api_client, workspace_partition, file_id)
+        assert set(ws_before) == {ws1, ws2}
+
+        # Update file metadata via PATCH
+        r = api_client.patch(
+            f"/indexer/partition/{workspace_partition}/file/{file_id}",
+            data={"metadata": '{"updated": true}'},
+        )
+        assert r.status_code == 200
+
+        # Workspace memberships must still be intact
+        ws_after = self._get_file_workspaces(api_client, workspace_partition, file_id)
+        assert set(ws_after) == {ws1, ws2}, (
+            f"Expected workspace memberships {{{ws1}, {ws2}}} after PATCH, got {set(ws_after)}"
+        )
+        assert ws3 not in ws_after