From 8e6f2c4ae368c628d4d25b083d9c616034767f77 Mon Sep 17 00:00:00 2001 From: Arnau Gris Date: Mon, 28 Sep 2026 12:32:39 +0200 Subject: [PATCH 1/5] Support optional connection IDs in SyncAsk --- .../src/hyperforge_nucliadb/sync/agent.py | 156 +++++++++++++- .../hyperforge_nucliadb/sync/config_driver.py | 17 +- .../src/hyperforge_nucliadb/sync/driver.py | 48 +++-- agents/nucliadb/tests/test_sync.py | 196 +++++++++++++++++- 4 files changed, 392 insertions(+), 25 deletions(-) diff --git a/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py b/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py index deff5a03..0348978e 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py +++ b/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py @@ -3,6 +3,7 @@ from typing import Dict, List, Optional, cast from uuid import uuid4 +from hyperforge import logger from hyperforge.configure import agent from hyperforge.interaction import ( Feedback, @@ -109,6 +110,135 @@ async def post_filter_resources_by_connection( memory: QuestionMemory, manager: Manager, resources: Dict[str, List[str]], + ) -> Dict[str, Dict[str, List[str]]]: + connections = self.get_connections(manager) + filtered_resources: Dict[str, Dict[str, List[str]]] = {} + for kb_source_id, resource_ids in resources.items(): + if connections[kb_source_id].config.connection_ids: + source_resources = await self._post_filter_configured_resources( + memory=memory, + manager=manager, + resources={kb_source_id: resource_ids}, + ) + else: + source_resources = await self._post_filter_hybrid_resources( + memory=memory, + manager=manager, + kb_source_id=kb_source_id, + resource_ids=resource_ids, + ) + filtered_resources.update(source_resources) + return filtered_resources + + async def _post_filter_hybrid_resources( + self, + memory: QuestionMemory, + manager: Manager, + kb_source_id: str, + resource_ids: List[str], + ) -> Dict[str, Dict[str, List[str]]]: + driver = self.sources[kb_source_id] + ndb = get_ndb_driver(manager, kb_source_id) + resource_objs: List[NucliaDBResource] = await asyncio.gather( + *[ + ndb.driver.get_resource_by_id( + kbid=ndb.config.kbid, + rid=resource_id, + query_params={"show": [ResourceProperties.ORIGIN.value]}, + ) + for resource_id in resource_ids + ] + ) + + public_resource_ids: List[str] = [] + connected_resources: Dict[str, List[str]] = {} + for resource_obj in resource_objs: + origin = resource_obj.origin + if ( + origin is None + or not origin.source_id + or not origin.source_id.startswith("sync_config_") + ): + public_resource_ids.append(resource_obj.id) + continue + + source_id = origin.source_id + sync_config_id = source_id.removeprefix("sync_config_") + if not sync_config_id or origin.sync_metadata is None: + logger.warning( + "Excluding synced resource with incomplete origin metadata", + extra={"resource_id": resource_obj.id, "source_id": source_id}, + ) + continue + connected_resources.setdefault(sync_config_id, []).append(resource_obj.id) + + resolved_resource_ids: List[str] = [] + for sync_config_id, connection_resource_ids in connected_resources.items(): + try: + await driver.resolve_sync_config(sync_config_id) + except Exception as exc: + logger.warning( + "Excluding resources for an unavailable sync configuration", + extra={ + "sync_config_id": sync_config_id, + "resource_count": len(connection_resource_ids), + "exception_type": type(exc).__name__, + }, + ) + continue + resolved_resource_ids.extend(connection_resource_ids) + + public_resource_id_set = set(public_resource_ids) + if not resolved_resource_ids: + return { + kb_source_id: { + "__hybrid__": [ + resource_id + for resource_id in resource_ids + if resource_id in public_resource_id_set + ] + } + } + + try: + connected_filtered_resources = await self._post_filter_configured_resources( + memory=memory, + manager=manager, + resources={kb_source_id: resolved_resource_ids}, + ) + except Exception as exc: + logger.warning( + "Excluding synced resources after authorization failed", + extra={ + "source_id": kb_source_id, + "exception_type": type(exc).__name__, + }, + ) + connected_filtered_resources = {} + + authorized_resource_ids = { + resource_id + for connection_resource_ids in connected_filtered_resources.get( + kb_source_id, {} + ).values() + for resource_id in connection_resource_ids + } + return { + kb_source_id: { + "__hybrid__": [ + resource_id + for resource_id in resource_ids + if resource_id in public_resource_id_set + or resource_id in authorized_resource_ids + ] + } + } + + async def _post_filter_configured_resources( + self, + memory: QuestionMemory, + manager: Manager, + resources: Dict[str, List[str]], ) -> Dict[str, Dict[str, List[str]]]: """We have a list of ARAG resources per KB source. Now we need to filter them by connection""" filtered_resources: Dict[str, Dict[str, List[str]]] = {} @@ -258,13 +388,19 @@ async def post_filter_resources_by_connection( raise Exception("Connection not found") credential = connection_credentials[inner_connection_id] + connection_resource_ids = connections_by_resource[connection_id] + connection_sync_metadata = { + resource_id: sync_metadata_by_resource[resource_id] + for resource_id in connection_resource_ids + if resource_id in sync_metadata_by_resource + } filtered_resource_ids = await self.validate_resources_by_connection( connection=source, credentials=credential, - resource_ids=resource_ids, + resource_ids=connection_resource_ids, connection_id=inner_connection_id, sync_config_id=connection_id, - sync_metadata_by_resource=sync_metadata_by_resource, + sync_metadata_by_resource=connection_sync_metadata, ) filtered_resources.setdefault(kb_source_id, {})[connection_id] = ( @@ -279,13 +415,20 @@ def enrich_catalog_filter( self, catalog_filter: Optional[CatalogFilterExpression] = None, ): + connection_ids = [ + connection_id + for source in self.config.sources + for connection_id in self.sources[source].sync_configs.keys() + ] + if not connection_ids: + return catalog_filter + if catalog_filter is None: catalog_filter = CatalogFilterExpression( resource=Or( operands=[ OriginSource(id=f"sync_config_{connection_id}") - for source in self.config.sources - for connection_id in self.sources[source].sync_configs.keys() + for connection_id in connection_ids ] ) ) @@ -296,10 +439,7 @@ def enrich_catalog_filter( Or( operands=[ OriginSource(id=f"sync_config_{connection_id}") - for source in self.config.sources - for connection_id in self.sources[ - source - ].sync_configs.keys() + for connection_id in connection_ids ] ), ] diff --git a/agents/nucliadb/src/hyperforge_nucliadb/sync/config_driver.py b/agents/nucliadb/src/hyperforge_nucliadb/sync/config_driver.py index a2395e23..7a6c6944 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb/sync/config_driver.py +++ b/agents/nucliadb/src/hyperforge_nucliadb/sync/config_driver.py @@ -2,6 +2,7 @@ from httpx import Timeout from hyperforge.driver import DriverConfig +from pydantic import Field, field_validator from pydantic.config import ConfigDict from hyperforge_nucliadb.driver_config import ( @@ -12,7 +13,21 @@ class SyncConnection(NucliaDBConnection): - connection_ids: list[str] + connection_ids: list[str] = Field( + default_factory=list, + description=( + "Sync configuration IDs to restrict searches to. When omitted or empty, " + "the entire knowledge box is searched and synced resources are authorized " + "individually." + ), + ) + + @field_validator("connection_ids") + @classmethod + def validate_connection_ids(cls, connection_ids: list[str]) -> list[str]: + if any(not connection_id.strip() for connection_id in connection_ids): + raise ValueError("connection_ids cannot contain blank values") + return connection_ids @property def kb_url(self) -> str: diff --git a/agents/nucliadb/src/hyperforge_nucliadb/sync/driver.py b/agents/nucliadb/src/hyperforge_nucliadb/sync/driver.py index 72a6fcbb..711f72e1 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb/sync/driver.py +++ b/agents/nucliadb/src/hyperforge_nucliadb/sync/driver.py @@ -90,29 +90,47 @@ class SyncDriver(NucliaDBDriver): async def init(cls, driver: Any) -> "SyncDriver": sync_driver = cast(SyncDriverConfig, driver) client = await sync_connect(sync_driver.config) - information = {} - sync_configs: Dict[str, list[str]] = {} - for sync_config_id in sync_driver.config.connection_ids: - response = await client.get(f"/sync_config/{sync_config_id}") - response.raise_for_status() - data = response.json() - connection_id = data["external_connection"]["id"] - response = await client.get(f"/external_connection/{connection_id}") - response.raise_for_status() - sync_configs.setdefault(sync_config_id, []).append(connection_id) - data = response.json() - information[connection_id] = ExternalConnectionOutput(**data) - return cls( + instance = cls( provider=sync_driver.provider, name=sync_driver.name, async_driver=client, - information=information, - sync_configs=sync_configs, + information={}, + sync_configs={}, config=sync_driver.config, driver=await connect(cast(NucliaDBConnection, sync_driver.config)), manager=await manager_connect(cast(NucliaDBConnection, sync_driver.config)), _synonyms=None, ) + for sync_config_id in sync_driver.config.connection_ids: + await instance.resolve_sync_config(sync_config_id) + return instance + + async def resolve_sync_config( + self, sync_config_id: str + ) -> ExternalConnectionOutput: + if ( + self.config.connection_ids + and sync_config_id not in self.config.connection_ids + ): + raise ValueError( + f"Sync configuration {sync_config_id} is not configured for this driver" + ) + + existing_connection_ids = self.sync_configs.get(sync_config_id, []) + if existing_connection_ids: + return self.information[existing_connection_ids[0]] + + UUID(sync_config_id) + response = await self.async_driver.get(f"/sync_config/{sync_config_id}") + response.raise_for_status() + connection_id = response.json()["external_connection"]["id"] + + response = await self.async_driver.get(f"/external_connection/{connection_id}") + response.raise_for_status() + information = ExternalConnectionOutput(**response.json()) + self.sync_configs[sync_config_id] = [connection_id] + self.information[connection_id] = information + return information async def get_oauth_url( self, diff --git a/agents/nucliadb/tests/test_sync.py b/agents/nucliadb/tests/test_sync.py index 04cfd305..09f47fe6 100644 --- a/agents/nucliadb/tests/test_sync.py +++ b/agents/nucliadb/tests/test_sync.py @@ -1,7 +1,8 @@ import json import os +from types import SimpleNamespace from typing import Any, cast -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest @@ -12,6 +13,12 @@ from hyperforge.memory.memory import EphemeralSessionMemory from hyperforge.minimal_fixtures import cassette_nua_key from hyperforge.pubsub import UserToAgentInteraction +from pydantic import ValidationError + +from hyperforge_nucliadb.sync.agent import SyncAskAgent +from hyperforge_nucliadb.sync.config import SyncAskAgentConfig +from hyperforge_nucliadb.sync.config_driver import SyncConnection +from hyperforge_nucliadb.sync.driver import SyncDriver NUA_KEY = os.environ.get( "NUA_KEY", @@ -77,6 +84,193 @@ EXTERNAL_CONNECTION_ID = "019cade7-64ee-7389-bec6-888c8ff8d604" +async def test_sync_connection_defaults_to_hybrid_mode(): + config = { + "url": "https://example.com", + "manager": "https://example.com", + "kbid": "kb", + "key": "key", + "description": "test", + "filters": [], + } + + assert SyncConnection(**config).connection_ids == [] + with pytest.raises(ValidationError): + SyncConnection(**config, connection_ids=[" "]) + + +async def test_sync_driver_lazily_resolves_and_caches_connection(): + config = SyncConnection( + url="https://example.com", + manager="https://example.com", + kbid="00000000-0000-0000-0000-000000000010", + key="key", + description="test", + filters=[], + ) + client = AsyncMock(spec=AsyncClient) + client.get.side_effect = [ + _mock_sync_response( + "GET", + f"https://example.com/sync_config/{SYNC_CONFIG_ID}", + {"external_connection": {"id": EXTERNAL_CONNECTION_ID}}, + ), + _mock_sync_response( + "GET", + f"https://example.com/external_connection/{EXTERNAL_CONNECTION_ID}", + { + "id": EXTERNAL_CONNECTION_ID, + "kb_id": "00000000-0000-0000-0000-000000000010", + "created_by": "00000000-0000-0000-0000-000000000001", + "created_at": "2024-01-01T00:00:00Z", + "updated_at": "2024-01-01T00:00:00Z", + "provider": "sharefile_oauth", + }, + ), + ] + driver = SyncDriver.model_construct( + provider="sync", + name="sync", + config=config, + async_driver=client, + information={}, + sync_configs={}, + driver=MagicMock(), + manager=MagicMock(), + _synonyms=None, + ) + + first = await driver.resolve_sync_config(SYNC_CONFIG_ID) + second = await driver.resolve_sync_config(SYNC_CONFIG_ID) + + assert first is second + assert first.provider.value == "sharefile_oauth" + assert client.get.await_count == 2 + + +async def test_sync_driver_does_not_expand_configured_allowlist(): + config = SyncConnection( + url="https://example.com", + manager="https://example.com", + kbid="00000000-0000-0000-0000-000000000010", + key="key", + description="test", + filters=[], + connection_ids=[SYNC_CONFIG_ID], + ) + client = AsyncMock(spec=AsyncClient) + driver = SyncDriver.model_construct( + provider="sync", + name="sync", + config=config, + async_driver=client, + information={}, + sync_configs={}, + driver=MagicMock(), + manager=MagicMock(), + _synonyms=None, + ) + + with pytest.raises(ValueError, match="not configured"): + await driver.resolve_sync_config(EXTERNAL_CONNECTION_ID) + client.get.assert_not_awaited() + + +async def test_hybrid_catalog_filter_is_preserved(): + agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) + agent.sources = { + "source": SimpleNamespace(sync_configs={}) # type: ignore[dict-item] + } + catalog_filter = MagicMock() + + assert agent.enrich_catalog_filter(catalog_filter) is catalog_filter + + +async def test_hybrid_resources_without_connections_skip_authorization(): + agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) + driver = SimpleNamespace(config=SimpleNamespace(connection_ids=[])) + agent.sources = {"source": driver} # type: ignore[assignment] + resource = SimpleNamespace(id="public-resource", origin=None) + ndb = SimpleNamespace( + config=SimpleNamespace(kbid="kb"), + driver=SimpleNamespace(get_resource_by_id=AsyncMock(return_value=resource)), + ) + configured_filter = AsyncMock() + agent._post_filter_configured_resources = configured_filter + + with patch("hyperforge_nucliadb.sync.agent.get_ndb_driver", return_value=ndb): + result = await agent._post_filter_hybrid_resources( + memory=MagicMock(), + manager=MagicMock(), + kb_source_id="source", + resource_ids=["public-resource"], + ) + + assert result == {"source": {"__hybrid__": ["public-resource"]}} + configured_filter.assert_not_awaited() + + +async def test_hybrid_resources_authorize_only_valid_sync_origins(): + agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) + driver = SimpleNamespace( + config=SimpleNamespace(connection_ids=[]), + resolve_sync_config=AsyncMock(), + ) + agent.sources = {"source": driver} # type: ignore[assignment] + resources = [ + SimpleNamespace(id="public-resource", origin=None), + SimpleNamespace( + id="connected-resource", + origin=SimpleNamespace( + source_id=f"sync_config_{SYNC_CONFIG_ID}", + sync_metadata=SimpleNamespace(), + ), + ), + SimpleNamespace( + id="missing-metadata", + origin=SimpleNamespace( + source_id=f"sync_config_{SYNC_CONFIG_ID}", sync_metadata=None + ), + ), + SimpleNamespace( + id="invalid-connection", + origin=SimpleNamespace( + source_id="sync_config_not-a-uuid", + sync_metadata=SimpleNamespace(), + ), + ), + ] + ndb = SimpleNamespace( + config=SimpleNamespace(kbid="kb"), + driver=SimpleNamespace(get_resource_by_id=AsyncMock(side_effect=resources)), + ) + + async def resolve_sync_config(sync_config_id: str): + if sync_config_id == "not-a-uuid": + raise ValueError("invalid sync config") + return SimpleNamespace() + + driver.resolve_sync_config.side_effect = resolve_sync_config + agent._post_filter_configured_resources = AsyncMock( + return_value={"source": {SYNC_CONFIG_ID: ["connected-resource"]}} + ) + + with patch("hyperforge_nucliadb.sync.agent.get_ndb_driver", return_value=ndb): + result = await agent._post_filter_hybrid_resources( + memory=MagicMock(), + manager=MagicMock(), + kb_source_id="source", + resource_ids=[resource.id for resource in resources], + ) + + assert result == { + "source": { + "__hybrid__": ["public-resource", "connected-resource"], + } + } + agent._post_filter_configured_resources.assert_awaited_once() + + def _mock_sync_response( method: str, url: str, payload: dict[str, Any] ) -> httpx.Response: From 4abfd744a4d37277821e5308c6d82733c207b529 Mon Sep 17 00:00:00 2001 From: Arnau Gris Date: Mon, 28 Sep 2026 12:43:09 +0200 Subject: [PATCH 2/5] Fix connection ID handling in SyncAskAgent and update tests for catalog filter preservation --- .../src/hyperforge_nucliadb/sync/agent.py | 10 +++++++--- agents/nucliadb/tests/test_sync.py | 20 +++++++++++++++++-- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py b/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py index 0348978e..ff087694 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py +++ b/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py @@ -415,13 +415,17 @@ def enrich_catalog_filter( self, catalog_filter: Optional[CatalogFilterExpression] = None, ): + if any( + not self.sources[source].config.connection_ids + for source in self.config.sources + ): + return catalog_filter + connection_ids = [ connection_id for source in self.config.sources - for connection_id in self.sources[source].sync_configs.keys() + for connection_id in self.sources[source].config.connection_ids ] - if not connection_ids: - return catalog_filter if catalog_filter is None: catalog_filter = CatalogFilterExpression( diff --git a/agents/nucliadb/tests/test_sync.py b/agents/nucliadb/tests/test_sync.py index 09f47fe6..da374f7f 100644 --- a/agents/nucliadb/tests/test_sync.py +++ b/agents/nucliadb/tests/test_sync.py @@ -176,16 +176,32 @@ async def test_sync_driver_does_not_expand_configured_allowlist(): client.get.assert_not_awaited() -async def test_hybrid_catalog_filter_is_preserved(): +async def test_hybrid_catalog_filter_is_preserved_after_connection_resolution(): agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) agent.sources = { - "source": SimpleNamespace(sync_configs={}) # type: ignore[dict-item] + "source": SimpleNamespace( + config=SimpleNamespace(connection_ids=[]), + sync_configs={SYNC_CONFIG_ID: [EXTERNAL_CONNECTION_ID]}, + ) # type: ignore[dict-item] } catalog_filter = MagicMock() assert agent.enrich_catalog_filter(catalog_filter) is catalog_filter +async def test_mixed_sources_preserve_hybrid_catalog_filter(): + agent = SyncAskAgent(SyncAskAgentConfig(sources=["hybrid", "restricted"])) + agent.sources = { + "hybrid": SimpleNamespace(config=SimpleNamespace(connection_ids=[])), + "restricted": SimpleNamespace( + config=SimpleNamespace(connection_ids=[SYNC_CONFIG_ID]) + ), + } # type: ignore[dict-item] + catalog_filter = MagicMock() + + assert agent.enrich_catalog_filter(catalog_filter) is catalog_filter + + async def test_hybrid_resources_without_connections_skip_authorization(): agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) driver = SimpleNamespace(config=SimpleNamespace(connection_ids=[])) From 9715b9b7bea8e17d6ac4b34e87f0823877a49f33 Mon Sep 17 00:00:00 2001 From: Arnau Gris Date: Mon, 28 Sep 2026 13:12:27 +0200 Subject: [PATCH 3/5] Update tests and agent logic for sync --- .../src/hyperforge_nucliadb/sync/agent.py | 64 ++++++--- agents/nucliadb/tests/test_sync.py | 134 ++++++++++++++++++ 2 files changed, 178 insertions(+), 20 deletions(-) diff --git a/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py b/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py index ff087694..1a07041b 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py +++ b/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py @@ -154,14 +154,22 @@ async def _post_filter_hybrid_resources( connected_resources: Dict[str, List[str]] = {} for resource_obj in resource_objs: origin = resource_obj.origin - if ( - origin is None - or not origin.source_id - or not origin.source_id.startswith("sync_config_") - ): + if origin is None: public_resource_ids.append(resource_obj.id) continue + if not origin.source_id or not origin.source_id.startswith( + "sync_config_" + ): + if origin.sync_metadata is None: + public_resource_ids.append(resource_obj.id) + else: + logger.warning( + "Excluding resource with inconsistent Sync origin metadata", + extra={"resource_id": resource_obj.id}, + ) + continue + source_id = origin.source_id sync_config_id = source_id.removeprefix("sync_config_") if not sync_config_id or origin.sync_metadata is None: @@ -173,6 +181,7 @@ async def _post_filter_hybrid_resources( connected_resources.setdefault(sync_config_id, []).append(resource_obj.id) resolved_resource_ids: List[str] = [] + resolved_connection_ids: set[str] = set() for sync_config_id, connection_resource_ids in connected_resources.items(): try: await driver.resolve_sync_config(sync_config_id) @@ -187,6 +196,7 @@ async def _post_filter_hybrid_resources( ) continue resolved_resource_ids.extend(connection_resource_ids) + resolved_connection_ids.add(sync_config_id) public_resource_id_set = set(public_resource_ids) if not resolved_resource_ids: @@ -205,6 +215,7 @@ async def _post_filter_hybrid_resources( memory=memory, manager=manager, resources={kb_source_id: resolved_resource_ids}, + allowed_connection_ids=resolved_connection_ids, ) except Exception as exc: logger.warning( @@ -239,13 +250,21 @@ async def _post_filter_configured_resources( memory: QuestionMemory, manager: Manager, resources: Dict[str, List[str]], + allowed_connection_ids: Optional[set[str]] = None, ) -> Dict[str, Dict[str, List[str]]]: """We have a list of ARAG resources per KB source. Now we need to filter them by connection""" filtered_resources: Dict[str, Dict[str, List[str]]] = {} - connections_by_resource: Dict[str, List[str]] = {} - sync_metadata_by_resource: Dict[str, SyncMetadata] = {} connections = self.get_connections(manager) for kb_source_id, resource_ids in resources.items(): + driver = self.sources[kb_source_id] + configured_connection_ids = ( + allowed_connection_ids + if allowed_connection_ids is not None + else set(driver.config.connection_ids) + ) + connections_by_resource: Dict[str, List[str]] = {} + sync_metadata_by_resource: Dict[str, SyncMetadata] = {} + # Get resource source connection for each resource ndb = get_ndb_driver(manager, kb_source_id) get_resource_ids_tasks = [] @@ -261,24 +280,28 @@ async def _post_filter_configured_resources( *get_resource_ids_tasks ) for resource_obj in resource_objs: + origin = resource_obj.origin if ( - resource_obj.origin is not None - and resource_obj.origin.source_id is not None + origin is None + or not origin.source_id + or not origin.source_id.startswith("sync_config_") + or origin.sync_metadata is None ): - source_id = resource_obj.origin.source_id.replace( - "sync_config_", "" - ) - connections_by_resource.setdefault(source_id, []).append( - resource_obj.id - ) - if resource_obj.origin.sync_metadata is not None: - sync_metadata_by_resource[resource_obj.id] = ( - resource_obj.origin.sync_metadata - ) + continue + + source_id = origin.source_id.removeprefix("sync_config_") + if source_id not in configured_connection_ids: + continue + + connections_by_resource.setdefault(source_id, []).append( + resource_obj.id + ) + sync_metadata_by_resource[resource_obj.id] = origin.sync_metadata # Get providers needed for the resources allocated needed_providers_ids = list(connections_by_resource.keys()) - driver = self.sources[kb_source_id] + if not needed_providers_ids: + continue creds_providers = {} for sync_config_id in needed_providers_ids: @@ -460,6 +483,7 @@ async def search_by_title( catalog_filter: Optional[CatalogFilterExpression] = None, **kwargs, ) -> Dict[str, List[str]]: + self.get_connections(manager) catalog_filter = self.enrich_catalog_filter(catalog_filter) resources = await super().search_by_title( memory=memory, diff --git a/agents/nucliadb/tests/test_sync.py b/agents/nucliadb/tests/test_sync.py index da374f7f..082fa0c7 100644 --- a/agents/nucliadb/tests/test_sync.py +++ b/agents/nucliadb/tests/test_sync.py @@ -15,6 +15,7 @@ from hyperforge.pubsub import UserToAgentInteraction from pydantic import ValidationError +from hyperforge_nucliadb.basic_ask_agent import BasicAskAgent from hyperforge_nucliadb.sync.agent import SyncAskAgent from hyperforge_nucliadb.sync.config import SyncAskAgentConfig from hyperforge_nucliadb.sync.config_driver import SyncConnection @@ -202,6 +203,108 @@ async def test_mixed_sources_preserve_hybrid_catalog_filter(): assert agent.enrich_catalog_filter(catalog_filter) is catalog_filter +async def test_search_by_title_initializes_sources_before_enriching_filter(): + agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) + driver = SimpleNamespace(config=SimpleNamespace(connection_ids=[])) + manager = SimpleNamespace(drivers=SimpleNamespace(get=MagicMock(return_value=driver))) + + with patch.object( + BasicAskAgent, + "search_by_title", + new=AsyncMock(return_value={"source": []}), + ): + result = await agent.search_by_title( + memory=MagicMock(), + manager=manager, # type: ignore[arg-type] + title="title", + ) + + assert result == {"source": []} + assert agent.sources == {"source": driver} + + +async def test_configured_source_ignores_resources_outside_allowlist(): + agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) + driver = SimpleNamespace( + config=SimpleNamespace(connection_ids=[SYNC_CONFIG_ID]), + sync_configs={SYNC_CONFIG_ID: [EXTERNAL_CONNECTION_ID]}, + information={}, + ) + agent.sources = {"source": driver} # type: ignore[assignment] + resources = [ + SimpleNamespace(id="public-resource", origin=None), + SimpleNamespace( + id="other-connection-resource", + origin=SimpleNamespace( + source_id="sync_config_00000000-0000-0000-0000-000000000001", + sync_metadata=SimpleNamespace(), + ), + ), + ] + ndb = SimpleNamespace( + config=SimpleNamespace(kbid="kb"), + driver=SimpleNamespace(get_resource_by_id=AsyncMock(side_effect=resources)), + ) + manager = SimpleNamespace(drivers=SimpleNamespace(get=MagicMock(return_value=driver))) + + with patch("hyperforge_nucliadb.sync.agent.get_ndb_driver", return_value=ndb): + result = await agent._post_filter_configured_resources( + memory=MagicMock(), + manager=manager, # type: ignore[arg-type] + resources={"source": [resource.id for resource in resources]}, + ) + + assert result == {} + + +async def test_hybrid_dynamic_connection_is_authorized(): + agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) + driver = SimpleNamespace( + config=SimpleNamespace(connection_ids=[]), + sync_configs={SYNC_CONFIG_ID: [EXTERNAL_CONNECTION_ID]}, + information={ + EXTERNAL_CONNECTION_ID: SimpleNamespace(provider="sharefile_oauth") + }, + validate_resources=AsyncMock(return_value=["connected-resource"]), + ) + resource = SimpleNamespace( + id="connected-resource", + origin=SimpleNamespace( + source_id=f"sync_config_{SYNC_CONFIG_ID}", + sync_metadata=SimpleNamespace(model_dump=MagicMock(return_value={})), + ), + ) + ndb = SimpleNamespace( + config=SimpleNamespace(kbid="kb"), + driver=SimpleNamespace(get_resource_by_id=AsyncMock(return_value=resource)), + ) + manager = SimpleNamespace(drivers=SimpleNamespace(get=MagicMock(return_value=driver))) + memory = MagicMock() + memory.get_session_id.return_value = "session" + memory.send_feedback = AsyncMock( + return_value=UserToAgentInteraction( + request_id="session", + response=json.dumps( + { + "existing_credentials": { + SYNC_CONFIG_ID: {EXTERNAL_CONNECTION_ID: "credential"} + } + } + ), + ) + ) + + with patch("hyperforge_nucliadb.sync.agent.get_ndb_driver", return_value=ndb): + result = await agent._post_filter_configured_resources( + memory=memory, + manager=manager, # type: ignore[arg-type] + resources={"source": [resource.id]}, + allowed_connection_ids={SYNC_CONFIG_ID}, + ) + + assert result == {"source": {SYNC_CONFIG_ID: ["connected-resource"]}} + + async def test_hybrid_resources_without_connections_skip_authorization(): agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) driver = SimpleNamespace(config=SimpleNamespace(connection_ids=[])) @@ -287,6 +390,37 @@ async def resolve_sync_config(sync_config_id: str): agent._post_filter_configured_resources.assert_awaited_once() +async def test_hybrid_resource_with_sync_metadata_and_invalid_source_is_excluded(): + agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) + driver = SimpleNamespace( + config=SimpleNamespace(connection_ids=[]), + resolve_sync_config=AsyncMock(), + ) + agent.sources = {"source": driver} # type: ignore[assignment] + resource = SimpleNamespace( + id="invalid-sync-resource", + origin=SimpleNamespace( + source_id=None, + sync_metadata=SimpleNamespace(), + ), + ) + ndb = SimpleNamespace( + config=SimpleNamespace(kbid="kb"), + driver=SimpleNamespace(get_resource_by_id=AsyncMock(return_value=resource)), + ) + + with patch("hyperforge_nucliadb.sync.agent.get_ndb_driver", return_value=ndb): + result = await agent._post_filter_hybrid_resources( + memory=MagicMock(), + manager=MagicMock(), + kb_source_id="source", + resource_ids=[resource.id], + ) + + assert result == {"source": {"__hybrid__": []}} + driver.resolve_sync_config.assert_not_awaited() + + def _mock_sync_response( method: str, url: str, payload: dict[str, Any] ) -> httpx.Response: From 6f9d7d516530701be647f0ce1458d567a7803d93 Mon Sep 17 00:00:00 2001 From: Arnau Gris Date: Mon, 28 Sep 2026 14:21:48 +0200 Subject: [PATCH 4/5] Fix lint --- agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py b/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py index 1a07041b..3cc0b42d 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py +++ b/agents/nucliadb/src/hyperforge_nucliadb/sync/agent.py @@ -158,9 +158,7 @@ async def _post_filter_hybrid_resources( public_resource_ids.append(resource_obj.id) continue - if not origin.source_id or not origin.source_id.startswith( - "sync_config_" - ): + if not origin.source_id or not origin.source_id.startswith("sync_config_"): if origin.sync_metadata is None: public_resource_ids.append(resource_obj.id) else: From 8672b22efc6044c15f1fbb58409a42e3f4ff3285 Mon Sep 17 00:00:00 2001 From: Arnau Gris Date: Mon, 28 Sep 2026 14:24:43 +0200 Subject: [PATCH 5/5] Fix lint --- agents/nucliadb/tests/test_sync.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/agents/nucliadb/tests/test_sync.py b/agents/nucliadb/tests/test_sync.py index 082fa0c7..ddc6cca3 100644 --- a/agents/nucliadb/tests/test_sync.py +++ b/agents/nucliadb/tests/test_sync.py @@ -206,7 +206,9 @@ async def test_mixed_sources_preserve_hybrid_catalog_filter(): async def test_search_by_title_initializes_sources_before_enriching_filter(): agent = SyncAskAgent(SyncAskAgentConfig(sources=["source"])) driver = SimpleNamespace(config=SimpleNamespace(connection_ids=[])) - manager = SimpleNamespace(drivers=SimpleNamespace(get=MagicMock(return_value=driver))) + manager = SimpleNamespace( + drivers=SimpleNamespace(get=MagicMock(return_value=driver)) + ) with patch.object( BasicAskAgent, @@ -245,7 +247,9 @@ async def test_configured_source_ignores_resources_outside_allowlist(): config=SimpleNamespace(kbid="kb"), driver=SimpleNamespace(get_resource_by_id=AsyncMock(side_effect=resources)), ) - manager = SimpleNamespace(drivers=SimpleNamespace(get=MagicMock(return_value=driver))) + manager = SimpleNamespace( + drivers=SimpleNamespace(get=MagicMock(return_value=driver)) + ) with patch("hyperforge_nucliadb.sync.agent.get_ndb_driver", return_value=ndb): result = await agent._post_filter_configured_resources( @@ -278,7 +282,9 @@ async def test_hybrid_dynamic_connection_is_authorized(): config=SimpleNamespace(kbid="kb"), driver=SimpleNamespace(get_resource_by_id=AsyncMock(return_value=resource)), ) - manager = SimpleNamespace(drivers=SimpleNamespace(get=MagicMock(return_value=driver))) + manager = SimpleNamespace( + drivers=SimpleNamespace(get=MagicMock(return_value=driver)) + ) memory = MagicMock() memory.get_session_id.return_value = "session" memory.send_feedback = AsyncMock(