github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-10-26 13:16:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(jina): add process-local request admission (#6489)

Browse Source 
* 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>
This commit is contained in:
creed 2026-10-11 08:41:22 +08:00 committed by GitHub
parent 3234563db8
commit 3461c12896
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 1069 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@ -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

51
backend/docs/CONFIGURATION.md
View File

@ -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.

31
backend/packages/harness/deerflow/community/jina_ai/AGENTS.md
View File

@ -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`.

128
backend/packages/harness/deerflow/community/jina_ai/jina_client.py
View File

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

143
backend/packages/harness/deerflow/community/jina_ai/request_admission.py Normal file
View File

@ -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

2
backend/packages/harness/deerflow/community/jina_ai/tools.py
View File

@ -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

752
backend/tests/test_jina_request_admission.py Normal file
View File

@ -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="<p>ok</p>"))
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() == "<p>ok</p>"
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)