From 3461c128967231fd3491a1af537e00d2a63c4bae Mon Sep 17 00:00:00 2001
From: creed <115765389+97Three@users.noreply.github.com>
Date: Sun, 11 Oct 2026 08:41:22 +0800
Subject: [PATCH] feat(jina): add process-local request admission (#6489)
* feat(jina): add process-local request admission
Signed-off-by: 97three <2212371308@qq.com>
* test(jina): cover closed admission owner loop
Signed-off-by: 97three <2212371308@qq.com>
* test(jina): retain admission across API redirects
Signed-off-by: 97three <2212371308@qq.com>
* fix(jina): drain redirect response cleanup before releasing admission
Signed-off-by: 97three <2212371308@qq.com>
---------
Signed-off-by: 97three <2212371308@qq.com>
---
README.md | 2 +
backend/docs/CONFIGURATION.md | 51 +-
.../deerflow/community/jina_ai/AGENTS.md | 31 +-
.../deerflow/community/jina_ai/jina_client.py | 128 ++-
.../community/jina_ai/request_admission.py | 143 ++++
.../deerflow/community/jina_ai/tools.py | 2 +-
backend/tests/test_jina_request_admission.py | 752 ++++++++++++++++++
7 files changed, 1069 insertions(+), 40 deletions(-)
create mode 100644 backend/packages/harness/deerflow/community/jina_ai/request_admission.py
create mode 100644 backend/tests/test_jina_request_admission.py
diff --git a/README.md b/README.md
index a5405b6621..d5fb98efca 100644
--- a/README.md
+++ b/README.md
@@ -184,6 +184,8 @@ streams simulate chunks from a non-streaming response, while no-tool streams sta
Jina also accepts an opt-in `max_response_bytes` tool setting (positive integer; omitted/null disables it). It stops oversized decoded responses before extraction, with no partial success or retry. This leaves the 4096-character output cap unchanged and does not bound HTTPX decompressor allocations or wire bandwidth; see [response budget](backend/docs/CONFIGURATION.md#jina-response-byte-budget).
+ Jina can also bound concurrent HTTP attempts and FIFO waiting with the opt-in `request_admission` tool extra. One immutable budget is shared across clients and agent threads in each process; replicas have independent budgets. See [Jina request admission](backend/docs/CONFIGURATION.md#jina-request-admission).
+
Run `make doctor` at any time to verify your setup and get actionable fix hints.
If you are opening a GitHub issue about a local setup or runtime problem, run
`make support-bundle`. The command prints reporter next steps, writes a
diff --git a/backend/docs/CONFIGURATION.md b/backend/docs/CONFIGURATION.md
index 1e32145c60..ef79ca613a 100644
--- a/backend/docs/CONFIGURATION.md
+++ b/backend/docs/CONFIGURATION.md
@@ -763,11 +763,60 @@ cannot cancel work already started by Jina. Enabling retries can send up to `1 +
and incur additional cost. Successful content and final `Error:` results retain
the existing contract.
+#### Jina request admission
+
+On the Jina `web_fetch` tool entry, `request_admission` is omitted/null by default
+(disabled). To enable it, supply all three fields as YAML numbers:
+
+```yaml
+request_admission:
+ max_concurrent_requests: 3
+ max_queue_size: 32
+ max_wait_seconds: 10
+```
+
+These are example values, not recommended or measured optimal limits. Concurrency
+must be a positive integer, queue capacity a nonnegative integer (zero rejects
+whenever all permits are held), and wait a finite positive number. Booleans,
+strings, missing/unknown fields and malformed values return `Error:` before client
+creation. The public tool arguments remain unchanged.
+
+Before the first enabled call, omitted/null settings bypass admission and do not
+initialize a policy. The first enabled call freezes one policy for the process.
+After that, identical and omitted/null settings reuse it, so a live config reload
+cannot silently bypass the active cap; conflicts return `restart required`.
+Restart to change or disable enabled settings. Lead/subagents, client instances,
+threads and event loops share the enabled budget. Each worker process or replica
+has its own budget; this is neither a distributed quota nor a per-user/account
+limit, and it does not control requests per minute.
+
+Each physical attempt holds a permit through response consumption and response
+stream close, including errors and cancellation. One HTTP client is reused across
+the logical crawl's retries; closing its idle connection pool happens after the
+final attempt and does not hold an attempt permit. Enabled admission streams even
+without a byte cap, preserving complete text and charset decoding. Response cleanup
+drains before release, so cancellation or a retry deadline may take longer to return
+if local close is slow; cancellation still wins if cleanup also fails. This cannot
+stop work already running remotely. Readability extraction and retry backoff hold no
+permit. Every retry joins admission again. Existing retry eligibility, Retry-After
+floors and decoded response byte caps apply.
+
+Waiting is bounded FIFO: full queues return terminal local
+`Error: Local Jina request admission rejected: queue full` without HTTP dispatch.
+Expired waits return `Error: Local Jina request admission rejected: wait expired`.
+Neither is described as a provider/network failure, fabricates a provider 429, or
+triggers a retry. Cancellation removes queued tickets. Waiting consumes the original
+retry budget when retries are enabled; the earlier of that deadline and
+`max_wait_seconds` applies and is checked again at dispatch. With retries disabled,
+only queue wait is bounded by admission, and the existing HTTP timeout semantics
+still apply.
+
#### Jina response byte budget
On the same Jina `web_fetch` tool entry, optionally set `max_response_bytes: 1048576`
(for example, 1 MiB). This uses existing tool configuration extras; no model-facing
-argument is added. Omitted or `null` preserves the buffered default. An enabled
+argument is added. Omitted or `null` preserves the buffered default unless request
+admission is enabled. An enabled
value must be a positive integer; booleans, strings, fractions, zero and negative
values return `Error:` before HTTP client creation or network activity.
diff --git a/backend/packages/harness/deerflow/community/jina_ai/AGENTS.md b/backend/packages/harness/deerflow/community/jina_ai/AGENTS.md
index de5e0a4312..8d73abe8e9 100644
--- a/backend/packages/harness/deerflow/community/jina_ai/AGENTS.md
+++ b/backend/packages/harness/deerflow/community/jina_ai/AGENTS.md
@@ -1,18 +1,19 @@
# Jina web fetch
-Retries stay provider-local; default `max_retries=0` sends once. One monotonic
-budget bounds requests/waits; cancellation propagates. Retry 502/503/504 and
-connection-establishment failures; 429 requires valid Retry-After. Parse ASCII
-integer seconds or HTTP dates (including obsolete forms); past dates floor at
-zero. Wait max(server floor, existing budget-capped jittered 0.5–4s backoff).
-Never reduce server floors; unfit waits return the HTTP error. Reset hints each
-attempt. Auth/payment errors stay terminal. Tests: `test_jina_retries.py` and
-`test_jina_retry_after.py` (offline, not hosted-provider validation).
+Retries default to zero; one monotonic budget covers attempts/waits. Retry 502/503/504 and connect failures; 429 requires valid Retry-After.
+Accept ASCII seconds or HTTP dates (also obsolete); past dates floor at zero. Wait max(server floor, budget-capped jittered 0.5–4s backoff); unfit floors
+return HTTP errors. Reset hints each attempt; auth/payment errors stay terminal.
-`max_response_bytes`: null/omitted keeps buffered POST; otherwise positive int
-(excluding bool), validated before client creation. Stream `aiter_bytes`, count
-content-decoded bytes before text decoding; exact limit passes, excess closes
-and returns a body-free terminal Error before extraction, for every status.
-Keep one retry loop/deadline and reset the counter per response. Never re-decode
-compression or mutate HTTPX internals. Decoder allocations/wire bytes are outside
-the cap. Offline transports: `tests/test_jina_response_limit.py`.
+`max_response_bytes`: null keeps buffered POST unless admission is enabled.
+Positive int, not bool; validate before client creation. Count decoded
+`aiter_bytes` before retaining chunks; exact limit passes. Excess closes and
+returns a body-free terminal Error before extraction. Reset per
+response; never re-decode compression. Decoder/wire bytes are uncapped.
+
+`request_admission.py` owns one immutable process policy. Locked FIFO reserves before loop-safe wakeup. Cancellation/expiry return grants;
+leases cover response stream cleanup, never retry backoff, idle client-pool close, or readability. Reuse one client across retries.
+
+`_CleanupStream` wraps every response via an async response hook, before HTTPX reads/closes redirects. It depends on httpx 0.28
+stream-close ordering; re-verify on upgrades. Drain stream/client close on cancellation, with cancellation taking precedence over cleanup errors. Null bypasses only before first enablement; afterward all callers reuse the frozen policy until restart.
+Config: `backend/docs/CONFIGURATION.md`. Tests:
+`tests/test_jina_{client,retries,retry_after,response_limit,request_admission}.py`.
diff --git a/backend/packages/harness/deerflow/community/jina_ai/jina_client.py b/backend/packages/harness/deerflow/community/jina_ai/jina_client.py
index a9f74035d2..57951823d0 100644
--- a/backend/packages/harness/deerflow/community/jina_ai/jina_client.py
+++ b/backend/packages/harness/deerflow/community/jina_ai/jina_client.py
@@ -5,11 +5,14 @@ import os
import random
import re
import time
+from contextlib import asynccontextmanager, nullcontext
from datetime import UTC, datetime
from email.utils import parsedate_to_datetime
import httpx
+from .request_admission import JinaAdmissionError, get_admission
+
logger = logging.getLogger(__name__)
_api_key_warned = False
@@ -60,9 +63,72 @@ def _retry_after(value: str | None) -> float | None:
return None
+class _CleanupStream(httpx.AsyncByteStream):
+ """Keep the lease until HTTPX's underlying close completes, even on cancel.
+
+ HTTPX closes automatically at EOF and marks a response closed before awaiting
+ its stream. Protect the stream close itself, including cancellation at EOF;
+ a second response.aclose() alone cannot finish an interrupted stream close.
+ """
+
+ def __init__(self, stream: httpx.AsyncByteStream):
+ self.stream = stream
+
+ def __aiter__(self):
+ return self.stream.__aiter__()
+
+ async def aclose(self):
+ await _finish_cleanup(self.stream.aclose())
+
+
+async def _finish_cleanup(close):
+ cleanup = asyncio.create_task(close)
+ task = asyncio.current_task()
+ # Cancellation may already be propagating when an async context manager
+ # starts cleanup. Remember it so a later close failure cannot mask it.
+ cancelled = bool(task is not None and task.cancelling())
+ try:
+ while not cleanup.done():
+ try:
+ await asyncio.shield(cleanup)
+ except asyncio.CancelledError:
+ cancelled = True
+ cleanup.result()
+ finally:
+ if cancelled:
+ raise asyncio.CancelledError
+
+
+async def _protect_response_cleanup(response: httpx.Response):
+ # Response hooks run before HTTPX consumes intermediate redirect responses.
+ response.stream = _CleanupStream(response.stream)
+
+
+@asynccontextmanager
+async def _admission_client(options):
+ # Keep one client across retries, matching the default transport behavior.
+ # Cleanup is still shielded so cancellation always propagates after close.
+ client = httpx.AsyncClient(**options)
+ try:
+ client.event_hooks["response"].append(_protect_response_cleanup)
+ yield client
+ finally:
+ await _finish_cleanup(client.aclose())
+
+
class JinaClient:
async def crawl(
- self, url: str, return_format: str = "html", timeout: int = 10, proxy: str | None = None, trust_env: bool = True, *, max_retries: int = 0, retry_budget_seconds: float = 30.0, max_response_bytes: int | None = None
+ self,
+ url: str,
+ return_format: str = "html",
+ timeout: int = 10,
+ proxy: str | None = None,
+ trust_env: bool = True,
+ *,
+ max_retries: int = 0,
+ retry_budget_seconds: float = 30.0,
+ max_response_bytes: int | None = None,
+ request_admission: dict | None = None,
) -> str:
"""Fetch with optional bounded retries; cancellation always propagates."""
global _api_key_warned
@@ -77,6 +143,7 @@ class JinaClient:
_api_key_warned = True
logger.warning("Jina API key is not set. Provide your own key to access a higher rate limit. See https://jina.ai/reader for more information.")
data = {"url": url}
+ deadline = None
try:
if max_response_bytes is not None and (isinstance(max_response_bytes, bool) or not isinstance(max_response_bytes, int) or max_response_bytes <= 0):
raise ValueError("max_response_bytes must be a positive integer or null")
@@ -85,37 +152,47 @@ class JinaClient:
if isinstance(retry_budget_seconds, bool) or not isinstance(retry_budget_seconds, (int, float)) or not math.isfinite(retry_budget_seconds) or retry_budget_seconds <= 0:
raise ValueError("retry_budget_seconds must be a finite positive number")
- # HTTPX timeouts are per network phase, so use an outer deadline to
- # bound the complete request sequence (including waits and cleanup).
+ admission = get_admission(request_admission)
+
+ # HTTPX timeouts are per network phase. Without admission, the outer
+ # deadline covers requests, waits, and cleanup. With admission,
+ # cleanup may outlast the retry budget; cancellation propagates after close.
deadline = asyncio.get_running_loop().time() + retry_budget_seconds if max_retries else None
- async with asyncio.timeout_at(deadline):
+ async with asyncio.timeout_at(None if admission else deadline):
client_kwargs: dict[str, object] = {"trust_env": trust_env, "follow_redirects": True}
if proxy:
client_kwargs["proxy"] = proxy
- async with httpx.AsyncClient(**client_kwargs) as client:
+ async with _admission_client(client_kwargs) if admission else httpx.AsyncClient(**client_kwargs) as client:
delay = 0.5
for attempt in range(max_retries + 1):
remaining = deadline - asyncio.get_running_loop().time() if deadline is not None else None
if remaining is not None and remaining <= 0:
raise TimeoutError
- request_timeout = min(timeout, remaining) if remaining is not None else timeout
server_floor = None
try:
- if max_response_bytes is None:
- response = await client.post("https://r.jina.ai/", headers=headers, json=data, timeout=request_timeout)
- response_text = response.text
- else:
- async with client.stream("POST", "https://r.jina.ai/", headers=headers, json=data, timeout=request_timeout) as response:
- content = bytearray()
- # aiter_bytes decodes Content-Encoding once. Check before
- # retaining each chunk; HTTPX decoder allocations are outside this cap.
- async for chunk in response.aiter_bytes():
- if len(content) + len(chunk) > max_response_bytes:
- return f"Error: Jina API response exceeds max_response_bytes ({max_response_bytes})"
- content.extend(chunk)
- # Match HTTPX text semantics without decompressing again or
- # mutating the response's private buffered-content state.
- response_text = content.decode(response.encoding or "utf-8", errors="replace")
+ async with admission.attempt(remaining) if admission else nullcontext() as lease:
+ # Admission consumes the original logical budget. Start the
+ # network timer only after acquisition so queue expiry stays local.
+ async with asyncio.timeout_at(deadline if admission else None):
+ if lease is not None:
+ lease.check_deadline()
+ remaining = deadline - asyncio.get_running_loop().time() if deadline is not None else None
+ if remaining is not None and remaining <= 0:
+ raise TimeoutError
+ request_timeout = min(timeout, remaining) if remaining is not None else timeout
+ if max_response_bytes is None and admission is None:
+ response = await client.post("https://r.jina.ai/", headers=headers, json=data, timeout=request_timeout)
+ response_text = response.text
+ else:
+ async with client.stream("POST", "https://r.jina.ai/", headers=headers, json=data, timeout=request_timeout) as response:
+ content = bytearray()
+ # Count decoded bytes before retaining a chunk; HTTPX
+ # decompressor allocations remain outside the cap.
+ async for chunk in response.aiter_bytes():
+ if max_response_bytes is not None and len(content) + len(chunk) > max_response_bytes:
+ return f"Error: Jina API response exceeds max_response_bytes ({max_response_bytes})"
+ content.extend(chunk)
+ response_text = content.decode(response.encoding or "utf-8", errors="replace")
except (httpx.ConnectError, httpx.ConnectTimeout):
if attempt == max_retries:
raise
@@ -145,10 +222,15 @@ class JinaClient:
wait = max(wait, server_floor)
if wait >= remaining:
return f"Error: Jina API returned status {response.status_code}: {response_text}"
- await asyncio.sleep(wait)
+ async with asyncio.timeout_at(deadline if admission else None):
+ await asyncio.sleep(wait)
delay = min(delay * 2, 4.0)
+ except JinaAdmissionError as e:
+ error_message = f"Local Jina request admission rejected: {e}"
+ logger.warning(error_message)
+ return f"Error: {error_message}"
except Exception as e:
- if isinstance(e, TimeoutError) and max_retries and asyncio.get_running_loop().time() >= deadline:
+ if isinstance(e, TimeoutError) and max_retries and deadline is not None and asyncio.get_running_loop().time() >= deadline:
error_message = "Request to Jina API failed: retry time budget exhausted"
else:
error_message = f"Request to Jina API failed: {type(e).__name__}: {e}"
diff --git a/backend/packages/harness/deerflow/community/jina_ai/request_admission.py b/backend/packages/harness/deerflow/community/jina_ai/request_admission.py
new file mode 100644
index 0000000000..21685fd5ba
--- /dev/null
+++ b/backend/packages/harness/deerflow/community/jina_ai/request_admission.py
@@ -0,0 +1,143 @@
+"""Opt-in Jina attempt admission, shared by threads/loops in this process.
+
+Tool extra ``request_admission`` is null (disabled before first enablement) or a
+mapping containing all three Policy fields. Once enabled, the policy applies to
+all later Jina calls in the process and is immutable until process restart.
+Tickets reserve capacity under a short threading lock; only their owning loop
+touches the wakeup future. No credentials, URLs, or caller identities are keys.
+"""
+
+import asyncio
+import math
+import threading
+import time
+from collections import deque
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+
+
+class JinaAdmissionError(Exception):
+ """Terminal local rejection; never a provider status or retryable I/O error."""
+
+
+@dataclass(frozen=True)
+class Policy:
+ max_concurrent_requests: int
+ max_queue_size: int
+ max_wait_seconds: float
+
+ @classmethod
+ def parse(cls, value: object) -> "Policy":
+ if not isinstance(value, dict) or set(value) != {"max_concurrent_requests", "max_queue_size", "max_wait_seconds"}:
+ raise ValueError("request_admission requires max_concurrent_requests, max_queue_size, max_wait_seconds")
+ for name, minimum in (("max_concurrent_requests", 1), ("max_queue_size", 0)):
+ if type(value[name]) is not int or value[name] < minimum:
+ raise ValueError(f"request_admission.{name} must be an integer >= {minimum}")
+ wait = value["max_wait_seconds"]
+ try:
+ valid_wait = type(wait) in (int, float) and math.isfinite(wait) and wait > 0
+ except OverflowError:
+ valid_wait = False
+ if not valid_wait:
+ raise ValueError("request_admission.max_wait_seconds must be finite and positive")
+ return cls(**value)
+
+
+@dataclass(eq=False)
+class _Ticket:
+ loop: asyncio.AbstractEventLoop
+ wakeup: asyncio.Future
+ deadline: float
+ state: str = "queued"
+
+ def check_deadline(self) -> None:
+ if time.monotonic() >= self.deadline:
+ raise JinaAdmissionError("wait expired")
+
+ def wake(self) -> None:
+ if not self.wakeup.done():
+ self.wakeup.set_result(None)
+
+
+class Admission:
+ def __init__(self, policy: Policy):
+ self.policy = policy
+ self._lock = threading.Lock()
+ self._queue: deque[_Ticket] = deque()
+ self._active = 0
+
+ def _drain(self) -> None:
+ """Called only under _lock; reserve before waking to prevent barging."""
+ while self._queue:
+ ticket = self._queue[0]
+ expired = time.monotonic() >= ticket.deadline
+ if not expired and self._active >= self.policy.max_concurrent_requests:
+ break
+ self._queue.popleft()
+ ticket.state = "expired" if expired else "granted"
+ if not expired:
+ self._active += 1
+ try:
+ ticket.loop.call_soon_threadsafe(ticket.wake)
+ except RuntimeError: # An owner loop was closed without draining tasks.
+ if ticket.state == "granted":
+ self._active -= 1
+ ticket.state = "released"
+
+ def _finish(self, ticket: _Ticket) -> None:
+ with self._lock:
+ if ticket.state == "queued":
+ self._queue.remove(ticket)
+ elif ticket.state == "granted":
+ self._active -= 1
+ ticket.state = "released"
+ self._drain()
+
+ @asynccontextmanager
+ async def attempt(self, remaining: float | None) -> AsyncIterator[_Ticket]:
+ loop = asyncio.get_running_loop()
+ wait = self.policy.max_wait_seconds if remaining is None else min(remaining, self.policy.max_wait_seconds)
+ ticket = _Ticket(loop, loop.create_future(), time.monotonic() + wait)
+ with self._lock:
+ self._drain()
+ if wait <= 0:
+ raise JinaAdmissionError("wait expired")
+ if not self._queue and self._active < self.policy.max_concurrent_requests:
+ ticket.state = "granted"
+ self._active += 1
+ elif len(self._queue) >= self.policy.max_queue_size:
+ raise JinaAdmissionError("queue full")
+ else:
+ self._queue.append(ticket)
+ try:
+ if ticket.state == "queued":
+ try:
+ async with asyncio.timeout(max(0, ticket.deadline - time.monotonic())):
+ await ticket.wakeup
+ except TimeoutError:
+ raise JinaAdmissionError("wait expired") from None
+ # A grant can race timeout/cancellation or sit in a stalled loop.
+ if ticket.state != "granted" or time.monotonic() >= ticket.deadline:
+ raise JinaAdmissionError("wait expired")
+ yield ticket
+ finally:
+ self._finish(ticket)
+
+
+_registry_lock = threading.Lock()
+_admission: Admission | None = None
+
+
+def get_admission(value: object) -> Admission | None:
+ """Null is disabled initially; after enablement every caller shares the policy."""
+ global _admission
+ with _registry_lock:
+ if value is None:
+ return _admission
+ policy = Policy.parse(value)
+ if _admission is None:
+ _admission = Admission(policy)
+ elif _admission.policy != policy:
+ raise ValueError("Jina request_admission changed; restart required")
+ return _admission
diff --git a/backend/packages/harness/deerflow/community/jina_ai/tools.py b/backend/packages/harness/deerflow/community/jina_ai/tools.py
index fd1e5c27df..1eb035195c 100644
--- a/backend/packages/harness/deerflow/community/jina_ai/tools.py
+++ b/backend/packages/harness/deerflow/community/jina_ai/tools.py
@@ -62,7 +62,7 @@ async def web_fetch_tool(url: str) -> str:
timeout = _coerce_timeout(config.model_extra.get("timeout"), timeout)
proxy = _coerce_proxy(config.model_extra.get("proxy"))
trust_env = _coerce_bool(config.model_extra.get("trust_env"), trust_env)
- request_options = {key: config.model_extra[key] for key in ("max_retries", "retry_budget_seconds", "max_response_bytes") if key in config.model_extra}
+ request_options = {key: config.model_extra[key] for key in ("max_retries", "retry_budget_seconds", "max_response_bytes", "request_admission") if key in config.model_extra}
html_content = await jina_client.crawl(url, return_format="html", timeout=timeout, proxy=proxy, trust_env=trust_env, **request_options)
if isinstance(html_content, str) and html_content.startswith("Error:"):
return html_content
diff --git a/backend/tests/test_jina_request_admission.py b/backend/tests/test_jina_request_admission.py
new file mode 100644
index 0000000000..2873e6793d
--- /dev/null
+++ b/backend/tests/test_jina_request_admission.py
@@ -0,0 +1,752 @@
+"""Offline Jina admission regressions through the real HTTPX transport boundary."""
+
+import asyncio
+import json
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+
+import httpx
+import pytest
+
+from deerflow.community.jina_ai import request_admission as admission_module
+from deerflow.community.jina_ai import tools
+from deerflow.community.jina_ai.jina_client import JinaClient
+
+pytestmark = pytest.mark.anyio
+POLICY = {"max_concurrent_requests": 1, "max_queue_size": 3, "max_wait_seconds": 5}
+
+
+@pytest.fixture
+def anyio_backend():
+ return "asyncio"
+
+
+@pytest.fixture(autouse=True)
+def isolated_policy(monkeypatch):
+ monkeypatch.setattr(admission_module, "_admission", None)
+ yield
+ budget = admission_module._admission
+ if budget is not None:
+ assert budget._active == 0
+ assert not budget._queue
+
+
+@pytest.fixture
+def install_transport(monkeypatch):
+ original = httpx.AsyncClient
+ monkeypatch.setenv("JINA_API_KEY", "dummy-test-key")
+
+ def install(handler):
+ monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: original(transport=httpx.MockTransport(handler), trust_env=False))
+
+ return install
+
+
+async def test_tool_rejects_full_queue_before_dispatch(monkeypatch, install_transport):
+ entered, release = asyncio.Event(), asyncio.Event()
+ calls = 0
+
+ async def handle(request):
+ nonlocal calls
+ calls += 1
+ entered.set()
+ if calls == 1:
+ await release.wait()
+ return httpx.Response(401, text="offline status")
+
+ install_transport(handle)
+ config = SimpleNamespace(model_extra={"request_admission": {**POLICY, "max_queue_size": 0}})
+ monkeypatch.setattr(tools, "get_app_config", lambda: SimpleNamespace(get_tool_config=lambda name: config))
+ first = asyncio.create_task(tools.web_fetch_tool.ainvoke("https://example.com/first"))
+ try:
+ await asyncio.wait_for(entered.wait(), 2)
+ result = await tools.web_fetch_tool.ainvoke("https://example.com/second")
+ assert result == "Error: Local Jina request admission rejected: queue full"
+ assert calls == 1
+ assert set(tools.web_fetch_tool.args) == {"url"}
+ finally:
+ release.set()
+ await first
+
+
+async def until(predicate):
+ async with asyncio.timeout(3):
+ while not predicate():
+ await asyncio.sleep(0)
+
+
+def queued(count):
+ budget = admission_module._admission
+ if budget is None:
+ return False
+ with budget._lock:
+ return len(budget._queue) == count
+
+
+def crawl(name="page", policy=POLICY, **kwargs):
+ return JinaClient().crawl(f"https://example.com/{name}", request_admission=policy, **kwargs)
+
+
+class Stream(httpx.AsyncByteStream):
+ def __init__(self, *, chunks=(b"ok",), read_error=None, close_error=None, hold_read=False, hold_close=False):
+ self.chunks = chunks
+ self.read_error = read_error
+ self.close_error = close_error
+ self.hold_read = hold_read
+ self.hold_close = hold_close
+ self.reading = asyncio.Event()
+ self.closing = asyncio.Event()
+ self.release_read = asyncio.Event()
+ self.release_close = asyncio.Event()
+ self.closed = False
+
+ async def __aiter__(self):
+ self.reading.set()
+ if self.hold_read:
+ await self.release_read.wait()
+ if self.read_error:
+ raise self.read_error
+ for chunk in self.chunks:
+ yield chunk
+
+ async def aclose(self):
+ self.closing.set()
+ if self.hold_close:
+ await self.release_close.wait()
+ self.closed = True
+ if self.close_error:
+ raise self.close_error
+
+
+@pytest.mark.parametrize("policy", [None, POLICY])
+@pytest.mark.parametrize("cap", [None, 2])
+async def test_default_concurrency_and_enabled_fifo(install_transport, policy, cap):
+ entered = []
+ releases = [asyncio.Event() for _ in range(8)]
+
+ async def handle(request):
+ number = int(json.loads(request.content)["url"].rsplit("/", 1)[1])
+ entered.append(number)
+ await releases[number].wait()
+ return httpx.Response(200, text="ok")
+
+ install_transport(handle)
+ options = None if policy is None else {**policy, "max_queue_size": 7}
+ tasks = []
+ try:
+ for i in range(8):
+ tasks.append(asyncio.create_task(crawl(str(i), options, max_response_bytes=cap)))
+ await until(lambda: len(entered) == i + 1 if policy is None else len(entered) == 1 and queued(i))
+ assert entered == (list(range(8)) if policy is None else [0])
+ for i in range(8):
+ releases[i].set()
+ assert await tasks[i] == "ok"
+ if policy is not None and i < 7:
+ await until(lambda: len(entered) == i + 2)
+ assert entered == list(range(i + 2))
+ finally:
+ for event in releases:
+ event.set()
+ await asyncio.gather(*tasks, return_exceptions=True)
+
+
+@pytest.mark.parametrize("cancel_index", [0, 1])
+async def test_queued_cancellation_removes_head_or_middle(install_transport, cancel_index):
+ stream = Stream(hold_read=True)
+ names = []
+
+ def handle(request):
+ names.append(json.loads(request.content)["url"].rsplit("/", 1)[1])
+ return httpx.Response(200, stream=stream) if len(names) == 1 else httpx.Response(200, text="ok")
+
+ install_transport(handle)
+ active = asyncio.create_task(crawl("active"))
+ await stream.reading.wait()
+ tasks = []
+ try:
+ for i in range(3):
+ tasks.append(asyncio.create_task(crawl(str(i))))
+ await until(lambda: queued(i + 1))
+ tasks[cancel_index].cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await tasks[cancel_index]
+ assert queued(2)
+ stream.release_read.set()
+ await active
+ await asyncio.gather(*tasks, return_exceptions=True)
+ assert names == ["active"] + [str(i) for i in range(3) if i != cancel_index]
+ finally:
+ stream.release_read.set()
+ await asyncio.gather(active, *tasks, return_exceptions=True)
+
+
+@pytest.mark.parametrize("retry_budget", [None, 0.02])
+async def test_wait_timeout_is_terminal_and_consumes_retry_budget(install_transport, retry_budget):
+ stream = Stream(hold_read=True)
+ calls = []
+ install_transport(lambda request: calls.append(request) or httpx.Response(200, stream=stream))
+ policy = {**POLICY, "max_wait_seconds": 0.02 if retry_budget is None else 5}
+ first = asyncio.create_task(crawl(policy=policy))
+ await stream.reading.wait()
+ try:
+ options = {} if retry_budget is None else {"max_retries": 2, "retry_budget_seconds": retry_budget}
+ result = await asyncio.wait_for(crawl(policy=policy, **options), 1)
+ assert result == "Error: Local Jina request admission rejected: wait expired"
+ assert len(calls) == 1
+ assert queued(0)
+ finally:
+ stream.release_read.set()
+ await first
+
+
+@pytest.mark.parametrize("cap", [None, 2])
+@pytest.mark.parametrize("failure", ["transport", "status", "read", "close", "byte_cap"])
+async def test_failures_release_lease(install_transport, cap, failure):
+ calls = 0
+ stream = Stream(
+ chunks=(b"oversize",) if failure == "byte_cap" else (b"ok",),
+ read_error=httpx.ReadError("read failed") if failure == "read" else None,
+ close_error=RuntimeError("close failed") if failure == "close" else None,
+ )
+
+ def handle(request):
+ nonlocal calls
+ calls += 1
+ if calls > 1:
+ return httpx.Response(200, text="ok")
+ if failure == "transport":
+ raise httpx.ConnectError("connect failed")
+ return httpx.Response(401 if failure == "status" else 200, stream=stream)
+
+ install_transport(handle)
+ result = await crawl(max_response_bytes=cap)
+ if failure == "byte_cap" and cap is None:
+ assert result == "oversize"
+ else:
+ assert result.startswith("Error:")
+ if failure != "transport":
+ assert stream.closed
+ assert await crawl(max_response_bytes=cap) == "ok"
+
+
+@pytest.mark.parametrize("at_close", [False, True])
+async def test_repeated_active_cancel_waits_for_cleanup(install_transport, at_close):
+ stream = Stream(hold_read=not at_close, hold_close=True)
+ calls = []
+
+ def handle(request):
+ calls.append(request)
+ return httpx.Response(200, stream=stream) if len(calls) == 1 else httpx.Response(200, text="ok")
+
+ install_transport(handle)
+ first = asyncio.create_task(crawl())
+ second = None
+ try:
+ await (stream.closing if at_close else stream.reading).wait()
+ first.cancel()
+ await stream.closing.wait()
+ second = asyncio.create_task(crawl())
+ await until(lambda: queued(1))
+ first.cancel()
+ await asyncio.sleep(0)
+ assert len(calls) == 1 and not first.done()
+ stream.release_close.set()
+ with pytest.raises(asyncio.CancelledError):
+ await first
+ assert stream.closed
+ assert await second == "ok"
+ finally:
+ stream.release_read.set()
+ stream.release_close.set()
+ await asyncio.gather(*[task for task in (first, second) if task], return_exceptions=True)
+
+
+@pytest.mark.parametrize(
+ "field,values",
+ [
+ ("max_concurrent_requests", [True, False, 0, -1, 1.5, "2", None]),
+ ("max_queue_size", [True, False, -1, 1.5, "2", None]),
+ ("max_wait_seconds", [True, False, 0, -1, float("inf"), float("nan"), "2", None, 10**400]),
+ ],
+)
+async def test_invalid_policy_before_client_activity(monkeypatch, field, values):
+ def unexpected(**kwargs):
+ pytest.fail("invalid policy created a client")
+
+ monkeypatch.setattr(httpx, "AsyncClient", unexpected)
+ for value in values:
+ result = await crawl(policy={**POLICY, field: value})
+ assert result.startswith("Error:") and field in result
+ assert admission_module._admission is None
+
+
+@pytest.mark.parametrize("value", [False, True, 1, [], {}, {**POLICY, "typo": 1}])
+async def test_malformed_policy_before_client_activity(monkeypatch, value):
+ monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: pytest.fail("invalid policy created a client"))
+ assert "request_admission requires" in await crawl(policy=value)
+
+
+async def test_enabled_policy_is_immutable_and_applies_to_later_null_calls(install_transport):
+ calls = []
+ stream = Stream(hold_read=True)
+
+ def handle(request):
+ name = json.loads(request.content)["url"].rsplit("/", 1)[1]
+ calls.append(name)
+ return httpx.Response(200, stream=stream) if name == "enabled" else httpx.Response(200, text="ok")
+
+ install_transport(handle)
+ assert await crawl("before", policy=None) == "ok"
+ assert admission_module._admission is None
+ active = asyncio.create_task(crawl("enabled"))
+ omitted = None
+ try:
+ await stream.reading.wait()
+ budget = admission_module._admission
+ omitted = asyncio.create_task(crawl("omitted", policy=None))
+ await until(lambda: queued(1))
+ assert calls == ["before", "enabled"]
+ assert "restart required" in await crawl(policy={**POLICY, "max_concurrent_requests": 2})
+ assert admission_module._admission is budget
+ stream.release_read.set()
+ assert await active == "ok"
+ assert await omitted == "ok"
+ assert calls == ["before", "enabled", "omitted"]
+ finally:
+ stream.release_read.set()
+ await asyncio.gather(*[task for task in (active, omitted) if task], return_exceptions=True)
+
+
+async def test_threads_and_event_loops_share_one_budget(install_transport):
+ entered = []
+ lock = threading.Lock()
+ release = threading.Event()
+ policy = {**POLICY, "max_queue_size": 7}
+
+ async def handle(request):
+ with lock:
+ entered.append(json.loads(request.content)["url"])
+ assert await asyncio.to_thread(release.wait, 3)
+ return httpx.Response(200, text="ok")
+
+ install_transport(handle)
+ with ThreadPoolExecutor(max_workers=8) as executor:
+ futures = [executor.submit(asyncio.run, crawl(str(i), policy)) for i in range(8)]
+ try:
+ await until(lambda: admission_module._admission is not None and queued(7))
+ with lock:
+ assert len(entered) == 1
+ release.set()
+ results = await asyncio.gather(*(asyncio.wrap_future(f) for f in futures))
+ assert results == ["ok"] * 8
+ finally:
+ release.set()
+
+
+async def test_grant_cancel_race_returns_reserved_permit(install_transport, monkeypatch):
+ stream = Stream(hold_read=True)
+ calls = []
+
+ def handle(request):
+ calls.append(request)
+ return httpx.Response(200, stream=stream) if len(calls) == 1 else httpx.Response(200, text="ok")
+
+ install_transport(handle)
+ first = asyncio.create_task(crawl())
+ await stream.reading.wait()
+ second = asyncio.create_task(crawl())
+ await until(lambda: queued(1))
+ original = admission_module._Ticket.wake
+
+ def cancel_at_grant(ticket):
+ second.cancel()
+ original(ticket)
+
+ monkeypatch.setattr(admission_module._Ticket, "wake", cancel_at_grant)
+ stream.release_read.set()
+ await first
+ with pytest.raises(asyncio.CancelledError):
+ await second
+ assert await crawl() == "ok"
+
+
+async def test_expired_grant_never_dispatches(install_transport, monkeypatch):
+ stream = Stream(hold_read=True)
+ calls = []
+
+ def handle(request):
+ calls.append(request)
+ return httpx.Response(200, stream=stream) if len(calls) == 1 else httpx.Response(200, text="ok")
+
+ install_transport(handle)
+ first = asyncio.create_task(crawl())
+ await stream.reading.wait()
+ second = asyncio.create_task(crawl())
+ await until(lambda: queued(1))
+ original = admission_module._Ticket.wake
+
+ def expire_at_grant(ticket):
+ ticket.deadline = 0
+ original(ticket)
+
+ monkeypatch.setattr(admission_module._Ticket, "wake", expire_at_grant)
+ stream.release_read.set()
+ await first
+ assert "wait expired" in await second
+ assert len(calls) == 1
+ assert await crawl() == "ok"
+
+
+@pytest.mark.parametrize("expired", [False, True])
+async def test_closed_owner_loop_returns_reserved_permit(expired):
+ budget = admission_module.get_admission(POLICY)
+ main_loop = asyncio.get_running_loop()
+ active = admission_module._Ticket(main_loop, main_loop.create_future(), float("inf"), state="granted")
+
+ def closed_loop_ticket():
+ loop = asyncio.new_event_loop()
+ try:
+ return admission_module._Ticket(loop, loop.create_future(), 0 if expired else float("inf"))
+ finally:
+ loop.close()
+
+ with ThreadPoolExecutor(max_workers=1) as executor:
+ ticket = await asyncio.wrap_future(executor.submit(closed_loop_ticket))
+
+ with budget._lock:
+ budget._active = 1
+ budget._queue.append(ticket)
+
+ budget._finish(active)
+ assert ticket.state == "released"
+ assert budget._active == 0
+ assert not budget._queue
+
+ # Finishing the orphan again must not decrement a granted or expired ticket.
+ budget._finish(ticket)
+ assert budget._active == 0
+ async with budget.attempt(None):
+ assert budget._active == 1
+ assert budget._active == 0
+
+
+async def test_retry_releases_during_backoff_and_rejoins_fifo(install_transport, monkeypatch):
+ from deerflow.community.jina_ai import jina_client
+
+ backoff, resume = asyncio.Event(), asyncio.Event()
+ holding = Stream(hold_read=True)
+ order = []
+ original_sleep = asyncio.sleep
+
+ async def sleep(delay):
+ if delay > 0:
+ backoff.set()
+ await resume.wait()
+ else:
+ await original_sleep(delay)
+
+ def handle(request):
+ name = json.loads(request.content)["url"].rsplit("/", 1)[1]
+ order.append(name)
+ if order == ["retry"]:
+ return httpx.Response(503, text="retry", headers={"Retry-After": "0"})
+ return httpx.Response(200, stream=holding) if name == "holding" else httpx.Response(200, text="ok")
+
+ monkeypatch.setattr(jina_client.asyncio, "sleep", sleep)
+ install_transport(handle)
+ retry = asyncio.create_task(crawl("retry", max_retries=1))
+ await backoff.wait()
+ active = asyncio.create_task(crawl("holding"))
+ await holding.reading.wait()
+ earlier = asyncio.create_task(crawl("earlier"))
+ await until(lambda: queued(1))
+ resume.set()
+ await until(lambda: queued(2))
+ holding.release_read.set()
+ assert await asyncio.gather(retry, active, earlier) == ["ok"] * 3
+ assert order == ["retry", "holding", "earlier", "retry"]
+
+
+async def test_readability_failure_holds_no_permit(install_transport, monkeypatch):
+ install_transport(lambda request: httpx.Response(200, text="ok
"))
+ config = SimpleNamespace(model_extra={"request_admission": POLICY})
+ monkeypatch.setattr(tools, "get_app_config", lambda: SimpleNamespace(get_tool_config=lambda name: config))
+
+ def extract(*args, **kwargs):
+ assert admission_module._admission._active == 0
+ raise RuntimeError("extraction failed")
+
+ monkeypatch.setattr(tools.readability_extractor, "extract_article", extract)
+ with pytest.raises(RuntimeError, match="extraction failed"):
+ await tools.web_fetch_tool.ainvoke("https://example.com")
+ assert await crawl() == "ok
"
+
+
+async def test_capacity_greater_than_one_and_nonempty_full_queue(install_transport):
+ streams = [Stream(hold_read=True) for _ in range(3)]
+ calls = []
+ policy = {**POLICY, "max_concurrent_requests": 3, "max_queue_size": 2}
+
+ def handle(request):
+ calls.append(request)
+ return httpx.Response(200, stream=streams[len(calls) - 1]) if len(calls) <= 3 else httpx.Response(200, text="ok")
+
+ install_transport(handle)
+ tasks = [asyncio.create_task(crawl(str(i), policy)) for i in range(5)]
+ try:
+ await until(lambda: queued(2))
+ assert len(calls) == 3
+ assert "queue full" in await crawl("full", policy, max_retries=2)
+ assert len(calls) == 3
+ for stream in streams:
+ stream.release_read.set()
+ assert await asyncio.gather(*tasks) == ["ok"] * 5
+ finally:
+ for stream in streams:
+ stream.release_read.set()
+ await asyncio.gather(*tasks, return_exceptions=True)
+
+
+async def test_enabled_retries_reuse_one_client(monkeypatch):
+ original_client = httpx.AsyncClient
+ clients = []
+ calls = 0
+
+ def handle(request):
+ nonlocal calls
+ calls += 1
+ if calls == 1:
+ return httpx.Response(503, text="retry", headers={"Retry-After": "0"})
+ return httpx.Response(200, text="ok")
+
+ def client(**kwargs):
+ result = original_client(transport=httpx.MockTransport(handle), trust_env=False)
+ clients.append(result)
+ return result
+
+ monkeypatch.setattr(httpx, "AsyncClient", client)
+ assert await crawl(max_retries=1, retry_budget_seconds=2) == "ok"
+ assert calls == 2
+ assert len(clients) == 1
+
+
+async def test_client_pool_close_does_not_hold_attempt_permit(monkeypatch):
+ original_client = httpx.AsyncClient
+ closing, release = asyncio.Event(), asyncio.Event()
+ calls = []
+ clients = []
+
+ class Client(original_client):
+ async def aclose(self):
+ if self is clients[0]:
+ closing.set()
+ await release.wait()
+ await super().aclose()
+
+ def handle(request):
+ calls.append(request)
+ return httpx.Response(200, text="ok")
+
+ def client(**kwargs):
+ result = Client(transport=httpx.MockTransport(handle), trust_env=False)
+ clients.append(result)
+ return result
+
+ monkeypatch.setattr(httpx, "AsyncClient", client)
+ first = asyncio.create_task(crawl())
+ try:
+ await closing.wait()
+ assert await crawl("second") == "ok"
+ assert len(calls) == 2
+ assert queued(0)
+ release.set()
+ assert await first == "ok"
+ finally:
+ release.set()
+ await asyncio.gather(first, return_exceptions=True)
+
+
+async def test_cancel_with_response_and_client_close_failures_preserves_cancel_and_permit(monkeypatch):
+ original_client = httpx.AsyncClient
+ stream = Stream(hold_read=True, close_error=RuntimeError("response close failed"))
+ clients = []
+ calls = []
+
+ class Client(original_client):
+ async def aclose(self):
+ await super().aclose()
+ if self is clients[0]:
+ raise RuntimeError("client close failed")
+
+ def handle(request):
+ calls.append(request)
+ return httpx.Response(200, stream=stream) if len(calls) == 1 else httpx.Response(200, text="ok")
+
+ def client(**kwargs):
+ result = Client(transport=httpx.MockTransport(handle), trust_env=False)
+ clients.append(result)
+ return result
+
+ monkeypatch.setattr(httpx, "AsyncClient", client)
+ first = asyncio.create_task(crawl("cancelled"))
+ second = None
+ try:
+ await stream.reading.wait()
+ second = asyncio.create_task(crawl("next"))
+ await until(lambda: queued(1))
+ first.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await first
+ assert stream.closed
+ assert await second == "ok"
+ assert len(calls) == 2
+ finally:
+ stream.release_read.set()
+ await asyncio.gather(*[task for task in (first, second) if task], return_exceptions=True)
+
+
+async def test_successful_queue_wait_reduces_original_http_budget(install_transport, monkeypatch):
+ loop = asyncio.get_running_loop()
+ now = [loop.time()]
+ monkeypatch.setattr(loop, "time", lambda: now[0])
+ holding = Stream(hold_read=True)
+ second_stream = Stream(hold_read=True)
+ calls = []
+
+ def handle(request):
+ calls.append(request)
+ if len(calls) == 1:
+ return httpx.Response(200, stream=holding)
+ assert request.extensions["timeout"]["read"] == pytest.approx(0.25)
+ return httpx.Response(200, stream=second_stream)
+
+ install_transport(handle)
+ first = asyncio.create_task(crawl())
+ await holding.reading.wait()
+ second = asyncio.create_task(crawl(max_retries=1, retry_budget_seconds=0.5))
+ try:
+ await until(lambda: queued(1))
+ now[0] += 0.25
+ holding.release_read.set()
+ await first
+ await second_stream.reading.wait()
+ now[0] += 0.3
+ assert "retry time budget exhausted" in await second
+ assert second_stream.closed
+ assert len(calls) == 2
+ finally:
+ holding.release_read.set()
+ second_stream.release_read.set()
+ await asyncio.gather(first, second, return_exceptions=True)
+
+
+@pytest.mark.parametrize("limit", [1, 2])
+async def test_admission_preserves_decoded_byte_cap(install_transport, limit):
+ import gzip
+
+ stream = Stream(chunks=(gzip.compress(b"ok"),))
+ install_transport(lambda request: httpx.Response(200, stream=stream, headers={"Content-Encoding": "gzip"}))
+ result = await crawl(max_response_bytes=limit)
+ assert ("max_response_bytes" in result) if limit == 1 else result == "ok"
+ assert stream.closed
+
+
+async def test_client_close_failure_returns_permit(monkeypatch):
+ original_client = httpx.AsyncClient
+ clients = []
+
+ class Client(original_client):
+ async def aclose(self):
+ await super().aclose()
+ if self is clients[0]:
+ raise RuntimeError("client close failed")
+
+ def client(**kwargs):
+ result = Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, text="ok")), trust_env=False)
+ clients.append(result)
+ return result
+
+ monkeypatch.setattr(httpx, "AsyncClient", client)
+ assert "client close failed" in await crawl()
+ assert await crawl() == "ok"
+
+
+@pytest.mark.parametrize("max_retries", [0, 1])
+async def test_admitted_redirect_keeps_post_and_strips_cross_host_key(monkeypatch, max_retries):
+ requests = []
+
+ def handle(request):
+ assert json.loads(request.read()) == {"url": "https://example.com/page"}
+ assert admission_module._admission._active == 1
+ requests.append(request)
+ if request.url.host == "r.jina.ai":
+ return httpx.Response(307, headers={"Location": "https://redirect.example/final"})
+ return httpx.Response(200, text="Fetched page")
+
+ original = httpx.AsyncClient
+ monkeypatch.setenv("JINA_API_KEY", "dummy-test-key")
+ monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: original(transport=httpx.MockTransport(handle), **kwargs))
+
+ assert await crawl(max_retries=max_retries, trust_env=False) == "Fetched page"
+ assert [request.url.host for request in requests] == ["r.jina.ai", "redirect.example"]
+ assert all(request.method == "POST" for request in requests)
+ assert requests[0].headers["Authorization"] == "Bearer dummy-test-key"
+ assert "Authorization" not in requests[1].headers
+ assert admission_module._admission._active == 0
+
+
+@pytest.mark.parametrize("exit_reason", ["cancel", "deadline"])
+async def test_redirect_cleanup_finishes_before_next_permit(monkeypatch, exit_reason):
+ redirect = Stream(hold_close=True)
+ dispatched = []
+ loop = asyncio.get_running_loop()
+ original_time = loop.time
+ clock_offset = [0.0]
+ monkeypatch.setattr(loop, "time", lambda: original_time() + clock_offset[0])
+
+ def handle(request):
+ name = json.loads(request.content)["url"].rsplit("/", 1)[1]
+ dispatched.append(name)
+ if name == "first":
+ assert request.url.host == "r.jina.ai"
+ return httpx.Response(307, headers={"Location": "https://redirect.example/final"}, stream=redirect)
+ assert redirect.closed
+ return httpx.Response(200, text="ok")
+
+ original_client = httpx.AsyncClient
+ monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: original_client(transport=httpx.MockTransport(handle), **kwargs))
+ first = asyncio.create_task(crawl("first", trust_env=False, max_retries=int(exit_reason == "deadline"), retry_budget_seconds=1))
+ second = None
+ try:
+ await asyncio.wait_for(redirect.closing.wait(), 3)
+ second = asyncio.create_task(crawl("second", trust_env=False))
+ await until(lambda: queued(1))
+ if exit_reason == "cancel":
+ first.cancel()
+ else:
+ clock_offset[0] = 2.0
+ await until(lambda: first.cancelling() or first.done())
+ # Let cancellation and the HTTPX exception cleanup path run while the
+ # transport close is deliberately held at the event barrier.
+ for _ in range(3):
+ await asyncio.sleep(0)
+ assert not first.done()
+ assert not redirect.closed
+ assert dispatched == ["first"]
+ assert queued(1)
+ assert admission_module._admission._active == 1
+
+ redirect.release_close.set()
+ if exit_reason == "cancel":
+ with pytest.raises(asyncio.CancelledError):
+ await first
+ else:
+ assert "retry time budget exhausted" in await first
+ assert await second == "ok"
+ assert redirect.closed
+ assert dispatched == ["first", "second"]
+ assert admission_module._admission._active == 0
+ finally:
+ redirect.release_close.set()
+ await asyncio.gather(*(task for task in (first, second) if task is not None), return_exceptions=True)