CVE-2026-11466

@@ -133,22 +133,24 @@ def _reflect_get_subquery(self, query: str, intermediate_context: List[str]) ->
         )
         return self.llm.remove_think(chat_response.content), chat_response.total_tokens
 
-    def _retrieve_and_answer(self, query: str) -> Tuple[str, List[RetrievalResult], int]:
+    def _retrieve_and_answer(self, query: str, **kwargs) -> Tuple[str, List[RetrievalResult], int]:
         consume_tokens = 0
         if self.route_collection:
             selected_collections, n_token_route = self.collection_router.invoke(
-                query=query, dim=self.embedding_model.dimension
+                query=query, dim=self.embedding_model.dimension, **kwargs
             )
         else:
-            selected_collections = self.collection_router.all_collections
+            selected_collections = self.collection_router.filter_authorized_collection_names(
+                self.collection_router.all_collections, **kwargs
+            )
             n_token_route = 0
         consume_tokens += n_token_route
         all_retrieved_results = []
         for collection in selected_collections:
             log.color_print(f"<search> Search [{query}] in [{collection}]...  </search>\n")
             query_vector = self.embedding_model.embed_query(query)
             retrieved_results = self.vector_db.search_data(
-                collection=collection, vector=query_vector, query_text=query
+                collection=collection, vector=query_vector, query_text=query, **kwargs
             )
             all_retrieved_results.extend(retrieved_results)
         all_retrieved_results = deduplicate_results(all_retrieved_results)
@@ -245,7 +247,7 @@ def retrieve(self, query: str, **kwargs) -> Tuple[List[RetrievalResult], int, di
             log.color_print(f">> Iteration: {iter + 1}\n")
             followup_query, n_token0 = self._reflect_get_subquery(query, intermediate_contexts)
             intermediate_answer, retrieved_results, n_token1 = self._retrieve_and_answer(
-                followup_query
+                followup_query, **kwargs
             )
             supported_retrieved_results, n_token2 = self._get_supported_docs(
                 retrieved_results, followup_query, intermediate_answer

@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import List, Optional, Set, Tuple
 
 from deepsearcher.agent.base import BaseAgent
 from deepsearcher.llm.base import BaseLLM
@@ -39,6 +39,38 @@ def __init__(self, llm: BaseLLM, vector_db: BaseVectorDB, dim: int, **kwargs):
             for collection_info in self.vector_db.list_collections(dim=dim)
         ]
 
+    def _get_authorized_collection_set(self, **kwargs) -> Optional[Set[str]]:
+        authorized_collections = kwargs.get("authorized_collection_set")
+        if authorized_collections is None:
+            authorized_collections = kwargs.get("authorized_collections")
+        if authorized_collections is None:
+            return None
+        return set(authorized_collections)
+
+    def filter_authorized_collection_names(
+        self, collection_names: List[str], **kwargs
+    ) -> List[str]:
+        authorized_collection_set = self._get_authorized_collection_set(**kwargs)
+        if authorized_collection_set is None:
+            return collection_names
+
+        return [
+            collection_name
+            for collection_name in collection_names
+            if collection_name in authorized_collection_set
+        ]
+
+    def _filter_authorized_collections(self, collection_infos: List, **kwargs) -> List:
+        authorized_collection_set = self._get_authorized_collection_set(**kwargs)
+        if authorized_collection_set is None:
+            return collection_infos
+
+        return [
+            collection_info
+            for collection_info in collection_infos
+            if collection_info.collection_name in authorized_collection_set
+        ]
+
     def invoke(self, query: str, dim: int, **kwargs) -> Tuple[List[str], int]:
         """
         Determine which collections are relevant for the given query.
@@ -56,7 +88,9 @@ def invoke(self, query: str, dim: int, **kwargs) -> Tuple[List[str], int]:
                 - The token usage for the routing operation
         """
         consume_tokens = 0
-        collection_infos = self.vector_db.list_collections(dim=dim)
+        collection_infos = self._filter_authorized_collections(
+            self.vector_db.list_collections(dim=dim), **kwargs
+        )
         if len(collection_infos) == 0:
             log.warning(
                 "No collections found in the vector database. Please check the database connection."
@@ -83,6 +117,14 @@ def invoke(self, query: str, dim: int, **kwargs) -> Tuple[List[str], int]:
         )
         selected_collections = self.llm.literal_eval(chat_response.content)
         consume_tokens += chat_response.total_tokens
+        allowed_collection_names = {
+            collection_info.collection_name for collection_info in collection_infos
+        }
+        selected_collections = [
+            collection_name
+            for collection_name in selected_collections
+            if collection_name in allowed_collection_names
+        ]
 
         for collection_info in collection_infos:
             # If a collection description is not provided, use the query as the search query

@@ -117,14 +117,16 @@ def _generate_sub_queries(self, original_query: str) -> Tuple[List[str], int]:
         response_content = self.llm.remove_think(chat_response.content)
         return self.llm.literal_eval(response_content), chat_response.total_tokens
 
-    async def _search_chunks_from_vectordb(self, query: str, sub_queries: List[str]):
+    async def _search_chunks_from_vectordb(self, query: str, sub_queries: List[str], **kwargs):
         consume_tokens = 0
         if self.route_collection:
             selected_collections, n_token_route = self.collection_router.invoke(
-                query=query, dim=self.embedding_model.dimension
+                query=query, dim=self.embedding_model.dimension, **kwargs
             )
         else:
-            selected_collections = self.collection_router.all_collections
+            selected_collections = self.collection_router.filter_authorized_collection_names(
+                self.collection_router.all_collections, **kwargs
+            )
             n_token_route = 0
         consume_tokens += n_token_route
 
@@ -133,7 +135,7 @@ async def _search_chunks_from_vectordb(self, query: str, sub_queries: List[str])
         for collection in selected_collections:
             log.color_print(f"<search> Search [{query}] in [{collection}]...  </search>\n")
             retrieved_results = self.vector_db.search_data(
-                collection=collection, vector=query_vector, query_text=query
+                collection=collection, vector=query_vector, query_text=query, **kwargs
             )
             if not retrieved_results or len(retrieved_results) == 0:
                 log.color_print(
@@ -232,7 +234,7 @@ async def async_retrieve(
 
             # Create all search tasks
             search_tasks = [
-                self._search_chunks_from_vectordb(query, sub_gap_queries)
+                self._search_chunks_from_vectordb(query, sub_gap_queries, **kwargs)
                 for query in sub_gap_queries
             ]
             # Execute all tasks in parallel and wait for results

@@ -74,10 +74,12 @@ def retrieve(self, query: str, **kwargs) -> Tuple[List[RetrievalResult], int, di
         consume_tokens = 0
         if self.route_collection:
             selected_collections, n_token_route = self.collection_router.invoke(
-                query=query, dim=self.embedding_model.dimension
+                query=query, dim=self.embedding_model.dimension, **kwargs
             )
         else:
-            selected_collections = self.collection_router.all_collections
+            selected_collections = self.collection_router.filter_authorized_collection_names(
+                self.collection_router.all_collections, **kwargs
+            )
             n_token_route = 0
         consume_tokens += n_token_route
         all_retrieved_results = []
@@ -87,6 +89,7 @@ def retrieve(self, query: str, **kwargs) -> Tuple[List[RetrievalResult], int, di
                 vector=self.embedding_model.embed_query(query),
                 top_k=max(self.top_k // len(selected_collections), 1),
                 query_text=query,
+                **kwargs,
             )
             all_retrieved_results.extend(retrieval_res)
         all_retrieved_results = deduplicate_results(all_retrieved_results)
@@ -109,7 +112,7 @@ def query(self, query: str, **kwargs) -> Tuple[str, List[RetrievalResult], int]:
                 - A list of retrieved document results
                 - The total token usage
         """
-        all_retrieved_results, n_token_retrieval, _ = self.retrieve(query)
+        all_retrieved_results, n_token_retrieval, _ = self.retrieve(query, **kwargs)
         chunk_texts = []
         for chunk in all_retrieved_results:
             if self.text_window_splitter and "wider_text" in chunk.metadata:

@@ -88,7 +88,12 @@ def insert_data(self, documents: List[dict]):
         return [x.succeeded for x in result]
 
     def search_data(
-        self, collection: Optional[str], vector: List[float], top_k: int = 50
+        self,
+        collection: Optional[str],
+        vector: List[float],
+        top_k: int = 50,
+        *args,
+        **kwargs,
     ) -> List[RetrievalResult]:
         """Azure Cognitive Search implementation with compatibility for older SDK versions"""
         from azure.core.credentials import AzureKeyCredential

@@ -1,10 +1,10 @@
 import unittest
-from unittest.mock import MagicMock
+
 import numpy as np
 
-from deepsearcher.llm.base import BaseLLM, ChatResponse
 from deepsearcher.embedding.base import BaseEmbedding
-from deepsearcher.vector_db.base import BaseVectorDB, RetrievalResult, CollectionInfo
+from deepsearcher.llm.base import BaseLLM, ChatResponse
+from deepsearcher.vector_db.base import BaseVectorDB, CollectionInfo, RetrievalResult
 
 
 class MockLLM(BaseLLM):
@@ -43,7 +43,7 @@ def literal_eval(self, text):
             try:
                 import ast
                 return ast.literal_eval(text)
-            except:
+            except (SyntaxError, ValueError):
                 pass
                 
         return ["test_collection"]
@@ -101,6 +101,7 @@ def search_data(self, collection, vector, top_k=10, **kwargs):
         self.last_search_collection = collection
         self.last_search_vector = vector
         self.last_search_top_k = top_k
+        self.last_search_kwargs = kwargs
         
         return [
             RetrievalResult(
@@ -139,4 +140,4 @@ def setUp(self):
         """Set up test fixtures for agent tests."""
         self.llm = MockLLM()
         self.embedding_model = MockEmbedding(dimension=8)
-        self.vector_db = MockVectorDB() 
\ No newline at end of file
+        self.vector_db = MockVectorDB()

@@ -1,9 +1,8 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 
 from deepsearcher.agent import ChainOfRAG
-from deepsearcher.vector_db.base import RetrievalResult
 from deepsearcher.llm.base import ChatResponse
-
+from deepsearcher.vector_db.base import RetrievalResult
 from tests.agent.test_base import BaseAgentTest
 
 
@@ -83,6 +82,34 @@ def test_retrieve_and_answer(self):
         # Check the results
         self.assertEqual(answer, "Deep learning is a subset of machine learning that uses neural networks with multiple layers.")
         self.assertEqual(tokens, 15)  # 5 from collection_router + 10 from LLM
+
+    def test_retrieve_and_answer_forwards_authorization_context(self):
+        """Test vector DB search receives caller authorization context."""
+        query = "What is deep learning?"
+
+        self.chain_of_rag.collection_router.invoke = MagicMock(return_value=(["test_collection"], 5))
+        self.llm.chat = MagicMock(return_value=ChatResponse(
+            content="Deep learning is a subset of machine learning that uses neural networks with multiple layers.",
+            total_tokens=10
+        ))
+
+        answer, results, tokens = self.chain_of_rag._retrieve_and_answer(
+            query,
+            authorized_collection_set=["test_collection"],
+            tenant_id="tenant_a",
+        )
+
+        self.chain_of_rag.collection_router.invoke.assert_called_once_with(
+            query=query,
+            dim=self.embedding_model.dimension,
+            authorized_collection_set=["test_collection"],
+            tenant_id="tenant_a",
+        )
+        self.assertEqual(self.vector_db.last_search_kwargs["authorized_collection_set"], ["test_collection"])
+        self.assertEqual(self.vector_db.last_search_kwargs["tenant_id"], "tenant_a")
+        self.assertEqual(answer, "Deep learning is a subset of machine learning that uses neural networks with multiple layers.")
+        self.assertEqual(len(results), 3)
+        self.assertEqual(tokens, 15)
     
     def test_get_supported_docs(self):
         """Test the _get_supported_docs method."""
@@ -234,4 +261,4 @@ def test_format_retrieved_results(self):
 
 if __name__ == "__main__":
     import unittest
-    unittest.main() 
\ No newline at end of file
+    unittest.main()

@@ -3,7 +3,6 @@
 from deepsearcher.agent.collection_router import CollectionRouter
 from deepsearcher.llm.base import ChatResponse
 from deepsearcher.vector_db.base import CollectionInfo
-
 from tests.agent.test_base import BaseAgentTest
 
 
@@ -148,7 +147,47 @@ def test_invoke_with_no_description(self):
         self.assertEqual(set(selected_collections), {"with_desc", "no_desc"})
         self.assertEqual(tokens, 5)
 
+    def test_invoke_filters_authorized_collections(self):
+        """Test that routing only considers collections authorized for the caller."""
+        query = "quarterly revenue"
+
+        self.llm.chat = MagicMock(return_value=ChatResponse(
+            content='["science", "news"]',
+            total_tokens=10
+        ))
+
+        with patch('deepsearcher.utils.log.color_print'):
+            selected_collections, tokens = self.collection_router.invoke(
+                query,
+                dim=8,
+                authorized_collection_set=["news"],
+            )
+
+        self.assertEqual(selected_collections, ["news"])
+        self.assertEqual(tokens, 0)
+        self.llm.chat.assert_not_called()
+
+    def test_invoke_filters_llm_selected_unauthorized_collections(self):
+        """Test that LLM-selected collections are constrained by authorization."""
+        query = "quarterly revenue"
+
+        self.llm.chat = MagicMock(return_value=ChatResponse(
+            content='["science", "news"]',
+            total_tokens=10
+        ))
+
+        with patch('deepsearcher.utils.log.color_print'):
+            selected_collections, tokens = self.collection_router.invoke(
+                query,
+                dim=8,
+                authorized_collection_set=["books", "news"],
+            )
+
+        self.assertEqual(set(selected_collections), {"books", "news"})
+        self.assertNotIn("science", selected_collections)
+        self.assertEqual(tokens, 10)
+
 
 if __name__ == "__main__":
     import unittest
-    unittest.main() 
\ No newline at end of file
+    unittest.main()

@@ -1,9 +1,8 @@
-from unittest.mock import MagicMock, patch
 import asyncio
+from unittest.mock import MagicMock
 
 from deepsearcher.agent import DeepSearch
 from deepsearcher.vector_db.base import RetrievalResult
-
 from tests.agent.test_base import BaseAgentTest
 
 
@@ -74,6 +73,33 @@ def test_search_chunks_from_vectordb(self):
         # With our mock returning "YES" for RERANK_PROMPT, all chunks should be accepted
         self.assertEqual(len(results), 3)  # 3 mock results from MockVectorDB
         self.assertEqual(tokens, 35)  # 5 from collection_router + 10*3 from LLM calls for reranking
+
+    def test_search_chunks_from_vectordb_forwards_authorization_context(self):
+        """Test vector DB search receives caller authorization context."""
+        query = "What is deep learning?"
+        sub_queries = ["What is deep learning?"]
+
+        self.deep_search.collection_router.invoke = MagicMock(return_value=(["test_collection"], 5))
+
+        results, tokens = asyncio.run(
+            self.deep_search._search_chunks_from_vectordb(
+                query,
+                sub_queries,
+                authorized_collection_set=["test_collection"],
+                tenant_id="tenant_a",
+            )
+        )
+
+        self.deep_search.collection_router.invoke.assert_called_once_with(
+            query=query,
+            dim=self.embedding_model.dimension,
+            authorized_collection_set=["test_collection"],
+            tenant_id="tenant_a",
+        )
+        self.assertEqual(self.vector_db.last_search_kwargs["authorized_collection_set"], ["test_collection"])
+        self.assertEqual(self.vector_db.last_search_kwargs["tenant_id"], "tenant_a")
+        self.assertEqual(len(results), 3)
+        self.assertEqual(tokens, 35)
     
     def test_generate_gap_queries(self):
         """Test the _generate_gap_queries method."""
@@ -232,4 +258,4 @@ def test_format_chunk_texts(self):
 
 if __name__ == "__main__":
     import unittest
-    unittest.main() 
\ No newline at end of file
+    unittest.main()

@@ -2,7 +2,6 @@
 
 from deepsearcher.agent import NaiveRAG
 from deepsearcher.vector_db.base import RetrievalResult
-
 from tests.agent.test_base import BaseAgentTest
 
 
@@ -51,6 +50,29 @@ def test_retrieve(self):
         
         # Check token count
         self.assertEqual(tokens, 5)  # From our mocked collection_router.invoke
+
+    def test_retrieve_forwards_authorization_context(self):
+        """Test retrieve passes caller context to routing and vector search."""
+        query = "Test query"
+
+        self.naive_rag.collection_router.invoke = MagicMock(return_value=(["test_collection"], 5))
+
+        results, tokens, metadata = self.naive_rag.retrieve(
+            query,
+            authorized_collection_set=["test_collection"],
+            tenant_id="tenant_a",
+        )
+
+        self.naive_rag.collection_router.invoke.assert_called_once_with(
+            query=query,
+            dim=self.embedding_model.dimension,
+            authorized_collection_set=["test_collection"],
+            tenant_id="tenant_a",
+        )
+        self.assertEqual(self.vector_db.last_search_kwargs["authorized_collection_set"], ["test_collection"])
+        self.assertEqual(self.vector_db.last_search_kwargs["tenant_id"], "tenant_a")
+        self.assertEqual(tokens, 5)
+        self.assertEqual(len(results), 3)
     
     def test_retrieve_without_routing(self):
         """Test retrieve method with routing disabled."""
@@ -69,6 +91,28 @@ def test_retrieve_without_routing(self):
         
         # Check token count
         self.assertEqual(tokens, 0)  # No tokens used for routing
+
+    def test_retrieve_without_routing_filters_authorized_collections(self):
+        """Test disabled routing still respects authorized collections."""
+        self.vector_db._collections.append(
+            self.vector_db._collections[0].__class__(
+                collection_name="secret_collection",
+                description="Secret collection",
+            )
+        )
+        self.naive_rag.collection_router.all_collections = ["test_collection", "secret_collection"]
+        self.naive_rag.route_collection = False
+        query = "Test query without routing"
+
+        results, tokens, metadata = self.naive_rag.retrieve(
+            query,
+            authorized_collection_set=["test_collection"],
+        )
+
+        self.assertEqual(self.vector_db.last_search_collection, "test_collection")
+        self.assertEqual(self.vector_db.last_search_kwargs["authorized_collection_set"], ["test_collection"])
+        self.assertEqual(tokens, 0)
+        self.assertEqual(len(results), 3)
     
     def test_query(self):
         """Test the query method."""
@@ -127,4 +171,4 @@ def test_with_window_splitter_disabled(self):
 
 if __name__ == "__main__":
     import unittest
-    unittest.main() 
\ No newline at end of file
+    unittest.main()