diff --git a/agent/context_references.py b/agent/context_references.py index d77857584a779..fe63190e2c0d0 100644 --- a/agent/context_references.py +++ b/agent/context_references.py @@ -152,13 +152,24 @@ async def preprocess_context_references_async( blocks: list[str] = [] injected_tokens = 0 - for ref in refs: - warning, block = await _expand_reference( - ref, - cwd_path, - url_fetcher=url_fetcher, - allowed_root=allowed_root_path, + # Expand all references concurrently. Each _expand_reference is independent + # (no shared state during expansion) — a message with several @url: refs + # would otherwise pay one full web_extract round-trip per ref in series. + # gather preserves positional order, so we reassemble warnings/blocks in the + # original ref order exactly as the prior serial loop did; the token-budget + # check below is unchanged (it runs once, after all refs are expanded). + expanded = await asyncio.gather( + *( + _expand_reference( + ref, + cwd_path, + url_fetcher=url_fetcher, + allowed_root=allowed_root_path, + ) + for ref in refs ) + ) + for warning, block in expanded: if warning: warnings.append(warning) if block: diff --git a/tests/agent/test_context_refs_concurrent.py b/tests/agent/test_context_refs_concurrent.py new file mode 100644 index 0000000000000..a4dd64624b4ae --- /dev/null +++ b/tests/agent/test_context_refs_concurrent.py @@ -0,0 +1,53 @@ +"""Tests for concurrent @-reference expansion in context_references. + +RED before the refactor: test_refs_expand_concurrently asserts that N URL refs +(each a ~0.2s fetch) complete in roughly one fetch-time, not N×. On the serial +`for ref in refs: await` loop this FAILS (takes ~N×0.2s); after switching to +asyncio.gather it passes. The output-contract test guards that concurrency does +NOT change ordering, warnings, blocks, or token accounting. +""" +from __future__ import annotations + +import asyncio +import time + +import pytest + +from agent.context_references import preprocess_context_references_async + + +async def _slow_fetcher(url: str) -> str: + # Simulate a per-URL network fetch (web_extract round trip). + await asyncio.sleep(0.2) + return f"CONTENT[{url}]" + + +@pytest.mark.asyncio +async def test_refs_expand_concurrently(tmp_path): + # Three independent URL refs in one message. + msg = "see @url:https://a.example/x @url:https://b.example/y @url:https://c.example/z please" + t0 = time.perf_counter() + res = await preprocess_context_references_async( + msg, cwd=tmp_path, context_length=100_000, url_fetcher=_slow_fetcher, + ) + elapsed = time.perf_counter() - t0 + # Serial would be ~0.6s (3×0.2). Concurrent ~0.2s. Assert well under 2× one fetch. + assert elapsed < 0.4, f"expected concurrent (~0.2s), got {elapsed:.2f}s (serial?)" + # All three blocks present, in order. + assert res.expanded + body = res.message + assert body.index("a.example") < body.index("b.example") < body.index("c.example"), \ + "reference blocks must stay in original order" + + +@pytest.mark.asyncio +async def test_concurrent_preserves_output_contract(tmp_path): + """Concurrency must not change which blocks/warnings appear or their order.""" + msg = "@url:https://one.example/p @url:https://two.example/q" + res = await preprocess_context_references_async( + msg, cwd=tmp_path, context_length=100_000, url_fetcher=_slow_fetcher, + ) + assert "CONTENT[https://one.example/p]" in res.message + assert "CONTENT[https://two.example/q]" in res.message + assert res.message.index("one.example") < res.message.index("two.example") + assert res.injected_tokens > 0