typosaurus-sandbox/tests/test_research_scheduling.py

# retoor <retoor@molodetz.nl>

import asyncio
import unittest
from concurrent.futures import ThreadPoolExecutor
from typing import Any, AsyncIterator

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}"))


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()