fix: eph index key

Author: CalebMartinUiPath

packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py

@@ -696,8 +696,12 @@ def create_ephemeral_index(
         Returns:
             ContextGroundingIndex: The created index information.
         """
-        if folder_key is not None or folder_path is not None:
-            folder_key = self._resolve_folder_key(folder_key, folder_path)
+        # Resolve the folder key the same way retrieve_by_id does (falling back
+        # to the ambient folder context) so the index is created in the same
+        # folder scope it will later be retrieved from. Otherwise the index is
+        # created tenant-scoped but the GET sends x-uipath-folderkey, scoping the
+        # lookup to a folder the index doesn't live in -> 404 "Schema not found".
+        folder_key = self._resolve_folder_key(folder_key, folder_path)
         spec = self._create_ephemeral_spec(
             usage,
             attachments,
@@ -733,8 +737,12 @@ async def create_ephemeral_index_async(
         Returns:
             ContextGroundingIndex: The created index information.
         """
-        if folder_key is not None or folder_path is not None:
-            folder_key = self._resolve_folder_key(folder_key, folder_path)
+        # Resolve the folder key the same way retrieve_by_id does (falling back
+        # to the ambient folder context) so the index is created in the same
+        # folder scope it will later be retrieved from. Otherwise the index is
+        # created tenant-scoped but the GET sends x-uipath-folderkey, scoping the
+        # lookup to a folder the index doesn't live in -> 404 "Schema not found".
+        folder_key = self._resolve_folder_key(folder_key, folder_path)
         spec = self._create_ephemeral_spec(
             usage,
             attachments,
@@ -2262,11 +2270,13 @@ def _retrieve_by_id_spec(
         folder_key: Optional[str] = None,
         folder_path: Optional[str] = None,
     ) -> RequestSpec:
-        # Folder key not needed by retrieve by id. Adding it breaks ephemeral indexes
+        folder_key = self._resolve_folder_key(folder_key, folder_path)
+
         return RequestSpec(
             method="GET",
             endpoint=Endpoint(f"/ecs_/v2/indexes/{id}"),
             headers={
+                **header_folder(folder_key, None),
                 **header_job_key(),
             },
         )

packages/uipath-platform/tests/services/test_context_grounding_service.py

@@ -8,7 +8,6 @@
 from uipath.platform import UiPathApiConfig, UiPathExecutionContext
 from uipath.platform.common.constants import (
     ENV_JOB_KEY,
-    HEADER_FOLDER_KEY,
     HEADER_JOB_KEY,
     HEADER_USER_AGENT,
 )
@@ -3178,6 +3177,19 @@ def test_create_ephemeral_index(
     ) -> None:
         import uuid
 
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/orchestrator_/api/FoldersNavigation/GetFoldersForCurrentUser?searchText=test-folder-path&skip=0&take=20",
+            status_code=200,
+            json={
+                "PageItems": [
+                    {
+                        "Key": "test-folder-key",
+                        "FullyQualifiedName": "test-folder-path",
+                    }
+                ]
+            },
+        )
+
         httpx_mock.add_response(
             url=f"{base_url}{org}{tenant}/ecs_/v2/indexes/createephemeral",
             status_code=200,
@@ -3203,22 +3215,27 @@ def test_create_ephemeral_index(
         if sent_requests is None:
             raise Exception("No request was sent")
 
-        assert sent_requests[0].method == "POST"
-        assert (
-            sent_requests[0].url
-            == f"{base_url}{org}{tenant}/ecs_/v2/indexes/createephemeral"
+        create_request = next(
+            r
+            for r in sent_requests
+            if str(r.url).endswith("/ecs_/v2/indexes/createephemeral")
         )
+        assert create_request.method == "POST"
 
-        request_data = json.loads(sent_requests[0].content)
+        request_data = json.loads(create_request.content)
         assert request_data["usage"] == "DeepRAG"
         assert "dataSource" in request_data
         assert request_data["dataSource"]["attachments"] == [
             str(att) for att in attachment_ids
         ]
 
-        assert HEADER_USER_AGENT in sent_requests[0].headers
+        # Ambient folder context is resolved and scopes the create, matching how
+        # retrieve_by_id resolves the folder for the subsequent GET.
+        assert create_request.headers["x-uipath-folderkey"] == "test-folder-key"
+
+        assert HEADER_USER_AGENT in create_request.headers
         assert (
-            sent_requests[0].headers[HEADER_USER_AGENT]
+            create_request.headers[HEADER_USER_AGENT]
             == f"UiPath.Python.Sdk/UiPath.Python.Sdk.Activities.ContextGroundingService.create_ephemeral_index/{version}"
         )
 
@@ -3234,6 +3251,19 @@ async def test_create_ephemeral_index_async(
     ) -> None:
         import uuid
 
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/orchestrator_/api/FoldersNavigation/GetFoldersForCurrentUser?searchText=test-folder-path&skip=0&take=20",
+            status_code=200,
+            json={
+                "PageItems": [
+                    {
+                        "Key": "test-folder-key",
+                        "FullyQualifiedName": "test-folder-path",
+                    }
+                ]
+            },
+        )
+
         httpx_mock.add_response(
             url=f"{base_url}{org}{tenant}/ecs_/v2/indexes/createephemeral",
             status_code=200,
@@ -3259,22 +3289,27 @@ async def test_create_ephemeral_index_async(
         if sent_requests is None:
             raise Exception("No request was sent")
 
-        assert sent_requests[0].method == "POST"
-        assert (
-            sent_requests[0].url
-            == f"{base_url}{org}{tenant}/ecs_/v2/indexes/createephemeral"
+        create_request = next(
+            r
+            for r in sent_requests
+            if str(r.url).endswith("/ecs_/v2/indexes/createephemeral")
         )
+        assert create_request.method == "POST"
 
-        request_data = json.loads(sent_requests[0].content)
+        request_data = json.loads(create_request.content)
         assert request_data["usage"] == "DeepRAG"
         assert "dataSource" in request_data
         assert request_data["dataSource"]["attachments"] == [
             str(att) for att in attachment_ids
         ]
 
-        assert HEADER_USER_AGENT in sent_requests[0].headers
+        # Ambient folder context is resolved and scopes the create, matching how
+        # retrieve_by_id resolves the folder for the subsequent GET.
+        assert create_request.headers["x-uipath-folderkey"] == "test-folder-key"
+
+        assert HEADER_USER_AGENT in create_request.headers
         assert (
-            sent_requests[0].headers[HEADER_USER_AGENT]
+            create_request.headers[HEADER_USER_AGENT]
             == f"UiPath.Python.Sdk/UiPath.Python.Sdk.Activities.ContextGroundingService.create_ephemeral_index_async/{version}"
         )
 
@@ -3919,61 +3954,3 @@ def test_ingest_data_omits_job_key_header_when_env_unset(
 
         headers = mock_request.call_args[1]["headers"]
         assert HEADER_JOB_KEY not in headers
-
-
-class TestRetrieveByIdFolderScoping:
-    """retrieve_by_id must not send a folder header.
-
-    Indexes are looked up by id tenant-wide. Ephemeral indexes are created
-    tenant-scoped, so sending x-uipath-folderkey on the GET scopes the lookup to
-    a folder the index does not live in, producing a 404 "Schema not found".
-    Regression: ECS-1819.
-    """
-
-    def test_omits_folder_header_even_with_ambient_folder(
-        self,
-        service: ContextGroundingService,
-        monkeypatch: pytest.MonkeyPatch,
-    ) -> None:
-        monkeypatch.delenv(ENV_JOB_KEY, raising=False)
-        # Simulate running inside a folder context (e.g. UIPATH_FOLDER_KEY set).
-        service._folder_key = "ambient-folder-key"
-
-        with patch.object(service, "request") as mock_request:
-            mock_request.return_value = MagicMock()
-            service.retrieve_by_id("ephemeral-index-id")
-
-        headers = mock_request.call_args[1]["headers"]
-        assert HEADER_FOLDER_KEY not in headers
-
-    @pytest.mark.anyio
-    async def test_async_omits_folder_header_even_with_ambient_folder(
-        self,
-        service: ContextGroundingService,
-        monkeypatch: pytest.MonkeyPatch,
-    ) -> None:
-        monkeypatch.delenv(ENV_JOB_KEY, raising=False)
-        service._folder_key = "ambient-folder-key"
-
-        with patch.object(service, "request_async") as mock_request:
-            mock_request.return_value = MagicMock()
-            await service.retrieve_by_id_async("ephemeral-index-id")
-
-        headers = mock_request.call_args[1]["headers"]
-        assert HEADER_FOLDER_KEY not in headers
-
-    def test_still_carries_job_key(
-        self,
-        service: ContextGroundingService,
-        monkeypatch: pytest.MonkeyPatch,
-    ) -> None:
-        monkeypatch.setenv(ENV_JOB_KEY, "job-key-retrieve")
-        service._folder_key = "ambient-folder-key"
-
-        with patch.object(service, "request") as mock_request:
-            mock_request.return_value = MagicMock()
-            service.retrieve_by_id("ephemeral-index-id")
-
-        headers = mock_request.call_args[1]["headers"]
-        assert headers[HEADER_JOB_KEY] == "job-key-retrieve"
-        assert HEADER_FOLDER_KEY not in headers