# retoor import asyncio import unittest from concurrent.futures import ThreadPoolExecutor from typing import Any, AsyncIterator from unittest import mock from typosaurus_sandbox.research.cache import TTLCache from typosaurus_sandbox.research.client import RsearchClient from typosaurus_sandbox.research.config import ResearchConfig from typosaurus_sandbox.research.envelopes import ChatResponse, DescribeResponse, SearchResponse, SearchResult from typosaurus_sandbox.research.frontier import QueryFrontier from typosaurus_sandbox.research.pipeline import ResearchPipeline, WorkItem SEARCH_FIXTURE: dict[str, Any] = { "query": "subject", "source": "duckduckgo", "count": 1, "success": True, "error": None, "results": [ { "title": "Result", "url": "https://example.com/result", "description": "Description", "source": "example.com", "extra": {}, "index": 0, } ], } CHAT_FIXTURE: dict[str, Any] = { "response": "Answer", "prompt": "prompt", "json_mode": False, "cached": False, "error": None, "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost_usd": 0.0001}, } DESCRIBE_FIXTURE: dict[str, Any] = { "url": "https://example.com/page", "description": "Page description", "elapsed": 0.5, "timestamp": "2026-08-07T12:00:00Z", } IMAGES_FIXTURE: dict[str, Any] = { "query": "subject", "source": "wikimedia", "count": 1, "success": True, "error": None, "results": [ { "title": "Aurora borealis over Norway", "url": "https://example.com/image", "description": "Photograph of the aurora borealis.", "source": "wikimedia", "extra": {}, "index": 0, } ], } class _FakeClient: def __init__(self) -> None: self.config = ResearchConfig() self.search_calls: list[tuple[str, dict[str, Any]]] = [] self.chat_calls: list[str] = [] self.describe_calls: list[str] = [] self.delay_seconds = 0.0 self.active = 0 self.peak = 0 def search_cached( self, query: str, *, source: str | None = None, count: int | None = None, content: bool = False, type: str | None = None, deep: bool = False, ai: bool = False, cache: bool = True, ) -> None: return None def describe_cached(self, url: str) -> None: return None async def search( self, query: str, *, source: str | None = None, count: int | None = None, content: bool = False, type: str | None = None, deep: bool = False, ai: bool = False, cache: bool = True, ) -> SearchResponse: self.search_calls.append((query, {"type": type, "content": content, "count": count})) self.active += 1 self.peak = max(self.peak, self.active) try: if self.delay_seconds: await asyncio.sleep(self.delay_seconds) if type == "images": return SearchResponse.from_dict(IMAGES_FIXTURE) return SearchResponse.from_dict(SEARCH_FIXTURE) finally: self.active -= 1 async def chat(self, prompt: str, *, system: str | None = None, json_mode: bool = False, cache: bool = True) -> ChatResponse: self.chat_calls.append(prompt) return ChatResponse.from_dict(CHAT_FIXTURE) async def describe(self, url: str) -> DescribeResponse: self.describe_calls.append(url) return DescribeResponse.from_dict(DESCRIBE_FIXTURE) class TestFrontierScheduling(unittest.IsolatedAsyncioTestCase): async def test_configured_concurrency_bounds_pool_and_drains_frontier(self) -> None: config = ResearchConfig() self.assertEqual(config.max_concurrency, 8) frontier = QueryFrontier() for i in range(64): self.assertTrue(frontier.push_query(f"query {i}", "seed")) issued: list[str] = [] active = 0 peak = 0 async def worker() -> None: nonlocal active, peak active += 1 peak = max(peak, active) try: while True: query = frontier.pop_query() if query is None: return issued.append(query) await asyncio.sleep(0) finally: active -= 1 await asyncio.gather(*(worker() for _ in range(config.max_concurrency))) stats = frontier.snapshot() self.assertEqual(peak, config.max_concurrency) self.assertEqual(stats.queries_enqueued, 64) self.assertEqual(stats.queries_issued, 64) self.assertEqual(len(issued), 64) self.assertEqual(len(set(issued)), 64) self.assertEqual(frontier.pending_count(), 0) async def test_concurrent_pools_never_issue_same_query_twice(self) -> None: frontier = QueryFrontier() for i in range(50): frontier.push_query(f"variant {i}", "seed") issued: list[str] = [] async def pool(size: int) -> None: async def pull() -> None: while True: query = frontier.pop_query() if query is None: return issued.append(query) await asyncio.sleep(0) await asyncio.gather(*(pull() for _ in range(size))) await asyncio.gather(pool(4), pool(4)) stats = frontier.snapshot() self.assertEqual(stats.queries_enqueued, 50) self.assertEqual(stats.queries_issued, 50) self.assertEqual(len(issued), 50) self.assertEqual(len(set(issued)), 50) async def test_queries_enqueued_while_pool_running_are_drained(self) -> None: frontier = QueryFrontier() for i in range(8): frontier.push_query(f"early {i}", "seed") issued: list[str] = [] stop = asyncio.Event() async def worker() -> None: while not stop.is_set() or frontier.has_pending(): query = frontier.pop_query() if query is None: await asyncio.sleep(0) continue issued.append(query) await asyncio.sleep(0) workers = [asyncio.create_task(worker()) for _ in range(4)] await asyncio.sleep(0) for i in range(5): frontier.push_query(f"late {i}", "result") stop.set() await asyncio.gather(*workers) stats = frontier.snapshot() self.assertEqual(stats.queries_issued, 13) self.assertEqual(len(set(issued)), 13) self.assertEqual(frontier.pending_count(), 0) def test_snapshot_accounting_is_consistent(self) -> None: frontier = QueryFrontier("subject") self.assertFalse(frontier.push_query("subject")) self.assertTrue(frontier.push_query("second query")) self.assertTrue(frontier.register_url("https://example.com/a")) self.assertFalse(frontier.register_url("https://example.com/a")) self.assertTrue(frontier.register_content("body text")) self.assertFalse(frontier.register_content("body text")) stats = frontier.snapshot() self.assertEqual(stats.queries_generated, 3) self.assertEqual(stats.queries_enqueued, 2) self.assertEqual(stats.queries_duplicates_skipped, 1) self.assertEqual(stats.urls_seen, 1) self.assertEqual(stats.urls_duplicates_skipped, 1) self.assertEqual(stats.content_seen, 1) self.assertEqual(stats.content_duplicates_skipped, 1) class TestFrontierConcurrencyDedup(unittest.IsolatedAsyncioTestCase): async def test_duplicate_query_pushes_under_concurrency_enqueue_once(self) -> None: frontier = QueryFrontier() async def push() -> bool: return frontier.push_query("same query", "origin") results = await asyncio.gather(*(push() for _ in range(64))) stats = frontier.snapshot() self.assertEqual(results.count(True), 1) self.assertEqual(stats.queries_generated, 64) self.assertEqual(stats.queries_enqueued, 1) self.assertEqual(stats.queries_duplicates_skipped, 63) async def test_concurrent_url_registration_dedups(self) -> None: frontier = QueryFrontier() urls = [ "https://example.com/page", "https://EXAMPLE.com/page", "https://example.com/page/", ] * 21 + ["https://example.com/page"] def storm() -> list[bool]: with ThreadPoolExecutor(max_workers=16) as pool: return list(pool.map(frontier.register_url, urls)) results = await asyncio.to_thread(storm) stats = frontier.snapshot() self.assertEqual(results.count(True), 1) self.assertEqual(stats.urls_seen, 1) self.assertEqual(stats.urls_duplicates_skipped, 63) async def test_concurrent_content_registration_dedups(self) -> None: frontier = QueryFrontier() def storm() -> list[bool]: with ThreadPoolExecutor(max_workers=16) as pool: return list(pool.map(frontier.register_content, ["identical page body"] * 64)) results = await asyncio.to_thread(storm) stats = frontier.snapshot() self.assertEqual(results.count(True), 1) self.assertEqual(stats.content_seen, 1) self.assertEqual(stats.content_duplicates_skipped, 63) async def test_overlapping_results_registered_once_under_concurrency(self) -> None: frontier = QueryFrontier() for i in range(32): frontier.push_query(f"query {i}", "seed") issued: list[str] = [] async def worker() -> None: while True: query = frontier.pop_query() if query is None: return issued.append(query) frontier.register_result( SearchResult(title=query, url="https://example.com/shared", description="", source="s", extra={}) ) await asyncio.sleep(0) await asyncio.gather(*(worker() for _ in range(8))) stats = frontier.snapshot() self.assertEqual(len(issued), 32) self.assertEqual(stats.urls_seen, 1) self.assertEqual(stats.urls_duplicates_skipped, 31) async def test_same_content_different_urls_registered_once(self) -> None: frontier = QueryFrontier() async def register(index: int) -> bool: return frontier.register_result( SearchResult( title=f"title {index}", url=f"https://example.com/page/{index}", description="", source="s", content="identical page body", extra={}, ) ) results = await asyncio.gather(*(register(i) for i in range(16))) stats = frontier.snapshot() self.assertEqual(results.count(True), 16) self.assertEqual(stats.urls_seen, 16) self.assertEqual(stats.content_seen, 1) self.assertEqual(stats.content_duplicates_skipped, 15) class TestTTLCacheBehaviour(unittest.TestCase): def test_repeat_key_returns_cached_value(self) -> None: cache = TTLCache[str]("repeat", ttl_seconds=60.0) cache.set("key", "value") self.assertEqual(cache.get("key"), "value") self.assertIs(cache.get("key"), cache.get("key")) def test_unknown_key_returns_none(self) -> None: cache = TTLCache[str]("missing", ttl_seconds=60.0) self.assertIsNone(cache.get("absent")) def test_zero_ttl_boundary_immediately_expired(self) -> None: cache = TTLCache[str]("zero", ttl_seconds=0.0) cache.set("key", "value") self.assertIsNone(cache.get("key")) def test_negative_ttl_never_returns_value(self) -> None: cache = TTLCache[str]("negative", ttl_seconds=-1.0) cache.set("key", "value") self.assertIsNone(cache.get("key")) def test_fresh_entry_survives_within_ttl(self) -> None: cache = TTLCache[str]("fresh", ttl_seconds=60.0) cache.set("key", "value") self.assertEqual(cache.get("key"), "value") def test_set_overwrites_existing_entry(self) -> None: cache = TTLCache[str]("overwrite", ttl_seconds=60.0) cache.set("key", "first") cache.set("key", "second") self.assertEqual(cache.get("key"), "second") def test_clear_removes_all_entries(self) -> None: cache = TTLCache[str]("clear", ttl_seconds=60.0) for i in range(10): cache.set(f"key-{i}", f"value-{i}") cache.clear() for i in range(10): self.assertIsNone(cache.get(f"key-{i}")) def test_entry_expires_after_ttl_elapses(self) -> None: cache = TTLCache[str]("expiry", ttl_seconds=10.0) with mock.patch("typosaurus_sandbox.research.cache.time.monotonic", side_effect=[100.0, 100.0, 111.0]): cache.set("key", "value") self.assertEqual(cache.get("key"), "value") self.assertIsNone(cache.get("key")) def test_entry_expires_exactly_at_ttl_boundary(self) -> None: cache = TTLCache[str]("boundary", ttl_seconds=10.0) with mock.patch("typosaurus_sandbox.research.cache.time.monotonic", side_effect=[100.0, 100.0, 110.0]): cache.set("key", "value") self.assertEqual(cache.get("key"), "value") self.assertIsNone(cache.get("key")) class TestTTLCacheConcurrency(unittest.TestCase): def test_concurrent_distinct_keys_all_retrievable(self) -> None: cache = TTLCache[str]("concurrent", ttl_seconds=60.0) keys = [f"key-{i}" for i in range(256)] def worker(key: str) -> None: cache.set(key, key + "-value") self.assertEqual(cache.get(key), key + "-value") with ThreadPoolExecutor(max_workers=16) as pool: list(pool.map(worker, keys)) for key in keys: self.assertEqual(cache.get(key), key + "-value") def test_concurrent_same_key_sets_single_consistent_value(self) -> None: cache = TTLCache[str]("storm", ttl_seconds=60.0) values = [f"value-{i}" for i in range(128)] def worker(value: str) -> None: cache.set("shared", value) self.assertIn(cache.get("shared"), values) with ThreadPoolExecutor(max_workers=16) as pool: list(pool.map(worker, values)) self.assertIn(cache.get("shared"), values) self.assertEqual(len(cache._entries), 1) class TestPipelineSingleMechanism(unittest.IsolatedAsyncioTestCase): async def test_web_images_chat_describe_dispatch_through_single_request_mechanism(self) -> None: client = RsearchClient() recorded: list[tuple[Any, ...]] = [] def fake( method: str, path: str, params: dict[str, str] | None, payload: bytes | None, headers: dict[str, str] | None, timeout: float | None, ) -> tuple[int, dict[str, Any]]: recorded.append((method, path, params, payload, headers, timeout)) if path == "/search": return 200, SEARCH_FIXTURE if path == "/chat": return 200, CHAT_FIXTURE return 200, DESCRIBE_FIXTURE client._request = fake await client.search("subject") await client.search("subject", type="images") await client.chat("prompt") await client.describe("https://example.com/page") await client.describe_raw(b"\x89PNGpayload", mime_type="image/png") self.assertEqual(len(recorded), 5) search_calls = [call for call in recorded if call[1] == "/search"] self.assertEqual(len(search_calls), 2) self.assertEqual(search_calls[0][0], "GET") self.assertEqual(search_calls[1][0], "GET") web_params = dict(search_calls[0][2] or {}) images_params = dict(search_calls[1][2] or {}) self.assertEqual(web_params, {"query": "subject"}) self.assertEqual(images_params, {"query": "subject", "type": "images"}) self.assertEqual({k: v for k, v in images_params.items() if k != "type"}, web_params) chat_calls = [call for call in recorded if call[1] == "/chat"] describe_calls = [call for call in recorded if call[1] == "/describe"] self.assertEqual(len(chat_calls), 1) self.assertEqual(chat_calls[0][0], "POST") self.assertEqual(len(describe_calls), 2) self.assertEqual(describe_calls[0][0], "GET") self.assertEqual(describe_calls[1][0], "POST") class TestResearchPipeline(unittest.IsolatedAsyncioTestCase): async def test_pipeline_pool_size_bounded_by_configured_concurrency(self) -> None: client = _FakeClient() frontier = QueryFrontier() default_pipeline = ResearchPipeline(client, frontier) self.assertEqual(default_pipeline.pool_size, 8) self.assertEqual(default_pipeline.pool_size, ResearchConfig().max_concurrency) narrow_pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=3)) self.assertEqual(narrow_pipeline.pool_size, 3) floor_pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=0)) self.assertEqual(floor_pipeline.pool_size, 1) async def test_pipeline_drains_all_work_items_and_dedups(self) -> None: client = _FakeClient() frontier = QueryFrontier() async def items() -> AsyncIterator[WorkItem]: yield WorkItem("web", "subject") yield WorkItem("web", "subject") yield WorkItem("images", "subject") yield WorkItem("describe", "https://example.com/page") yield WorkItem("chat", "prompt") pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=4)) report = await pipeline.run(items()) self.assertEqual(report.requests_succeeded, 5) self.assertEqual(report.requests_failed, 0) self.assertEqual(len(report.outcomes), 5) self.assertTrue(all(outcome.success for outcome in report.outcomes)) endpoints = {outcome.item.kind: outcome.endpoint for outcome in report.outcomes} self.assertEqual(endpoints, {"web": "/search", "images": "/search", "describe": "/describe", "chat": "/chat"}) self.assertEqual(len(client.search_calls), 3) self.assertEqual(len(client.describe_calls), 1) self.assertEqual(len(client.chat_calls), 1) stats = frontier.snapshot() self.assertEqual(stats.urls_seen, 2) self.assertEqual(stats.urls_duplicates_skipped, 1) self.assertEqual(stats.queries_enqueued, 6) self.assertEqual(stats.queries_duplicates_skipped, 2) self.assertEqual(stats.content_seen, 2) async def test_pipeline_concurrency_bounded_by_pool_size(self) -> None: client = _FakeClient() client.delay_seconds = 0.02 frontier = QueryFrontier() async def items() -> AsyncIterator[WorkItem]: for i in range(12): yield WorkItem("web", f"subject {i}") pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=4)) report = await pipeline.run(items()) self.assertEqual(report.requests_succeeded, 12) self.assertGreaterEqual(client.peak, 2) self.assertLessEqual(client.peak, 4) if __name__ == "__main__": unittest.main()