Merge pull request #96963 from kshitijk4poor/refactor/fast-lane-consolidation

fix(compression): fast-lane follow-up — certification parity, worker-thread telemetry, caller-cap wire shape
This commit is contained in:
kshitij
2026-08-28 13:11:03 +05:30
committed by GitHub
2 changed files with 176 additions and 39 deletions
+79 -39
View File
@@ -484,23 +484,14 @@ def _aux_progress_active() -> bool:
@contextlib.contextmanager
def aux_progress_hook(hook):
"""Install *hook* as the current thread's aux forward-progress callback.
def _aux_thread_local_hook(local: threading.local, hook):
"""Install one thread-local hook callback and restore its prior value.
``hook=None`` is a no-op passthrough so callers can wire it
unconditionally. Re-entrant-safe: restores the previous hook on exit.
``hook=None`` (or any non-callable) is a no-op passthrough so callers can
wire it unconditionally. Re-entrant-safe: restores the previous hook on
exit. Shared by the forward-progress hook and the content-free timing
hooks — one save/restore implementation, three thread-local slots.
"""
prev = getattr(_aux_progress, "hook", None)
_aux_progress.hook = hook if callable(hook) else prev
try:
yield
finally:
_aux_progress.hook = prev
@contextlib.contextmanager
def _aux_timing_hook(local: threading.local, hook):
"""Install one content-free timing hook and restore its prior value."""
previous = getattr(local, "hook", None)
local.hook = hook if callable(hook) else previous
try:
@@ -509,6 +500,21 @@ def _aux_timing_hook(local: threading.local, hook):
local.hook = previous
@contextlib.contextmanager
def aux_progress_hook(hook):
"""Install *hook* as the current thread's aux forward-progress callback.
``hook=None`` is a no-op passthrough so callers can wire it
unconditionally. Re-entrant-safe: restores the previous hook on exit.
"""
with _aux_thread_local_hook(_aux_progress, hook):
yield
# Back-compat alias — the timing hooks were introduced with this name.
_aux_timing_hook = _aux_thread_local_hook
def _run_protected_sync_provider_call(
callback: Callable[[dict[str, Any]], Any],
kwargs: dict[str, Any],
@@ -540,14 +546,24 @@ def _run_protected_sync_provider_call(
raise AuxiliaryExplicitCancellation()
progress_hook = getattr(_aux_progress, "hook", None)
# Timing hooks ride along with the progress hook: _create_with_progress
# fires _notify_aux_dispatch/_notify_aux_provider_response from whichever
# thread runs the provider callback, so an owner-thread-only install would
# silently drop provider_dispatch_ms / time_to_first_progress_ms whenever
# the protected daemon path is taken.
dispatch_hook = getattr(_aux_dispatch, "hook", None)
provider_response_hook = getattr(_aux_provider_response, "hook", None)
provider_context = contextvars.copy_context()
done = threading.Event()
outcome: dict[str, Any] = {}
def _provider_worker() -> None:
try:
with aux_progress_hook(progress_hook), aux_interrupt_protection(
cancel_check=cancel_check
with (
aux_progress_hook(progress_hook),
_aux_thread_local_hook(_aux_dispatch, dispatch_hook),
_aux_thread_local_hook(_aux_provider_response, provider_response_hook),
aux_interrupt_protection(cancel_check=cancel_check),
):
outcome["result"] = callback(kwargs)
except BaseException as exc:
@@ -8457,6 +8473,38 @@ class CompressionFastLane(NamedTuple):
reasoning_config: Optional[Dict[str, Any]]
def _fast_lane_config_fields(
config: Dict[str, Any],
) -> tuple[str, str, bool, Optional[int]]:
"""Extract the fast-lane certification fields from one task config.
Returns ``(provider, model, non_reasoning, cap)``:
- ``provider``/``model``: normalized (stripped; provider lowercased).
- ``non_reasoning``: True only when ``reasoning_effort`` EXPLICITLY
disables thinking. Delegates to ``parse_reasoning_effort`` so every
spelling users can write in config.yaml (``none``, ``false``,
``disabled``, YAML boolean ``false``) certifies identically —
``_get_task_extra_body`` already uses the same parser to disable
reasoning, and the two predicates must not disagree. Empty/unset
(provider default) is NOT non-reasoning.
- ``cap``: positive int from ``max_output_tokens``, else None.
Booleans are config drift, never a cap (``int(True) == 1``).
"""
from hermes_constants import parse_reasoning_effort
provider = str(config.get("provider") or "").strip().lower()
model = str(config.get("model") or "").strip()
parsed_effort = parse_reasoning_effort(config.get("reasoning_effort"))
non_reasoning = parsed_effort is not None and parsed_effort.get("enabled") is False
raw_cap = config.get("max_output_tokens")
try:
cap = 0 if isinstance(raw_cap, bool) else int(raw_cap or 0)
except (TypeError, ValueError):
cap = 0
return provider, model, non_reasoning, (cap if cap > 0 else None)
def resolve_compression_fast_lane(
actual_provider: str,
actual_model: Optional[str],
@@ -8477,44 +8525,32 @@ def resolve_compression_fast_lane(
if route_config is not None
else _get_auxiliary_task_config("compression")
)
provider = str(requested_provider or config.get("provider") or "").strip().lower()
model = str(requested_model or config.get("model") or "").strip()
effort = str(config.get("reasoning_effort") or "").strip().lower()
cfg_provider, cfg_model, non_reasoning, cap = _fast_lane_config_fields(config)
provider = str(requested_provider or "").strip().lower() or cfg_provider
model = str(requested_model or "").strip() or cfg_model
explicit_route = provider not in {"", "auto"} and model.lower() not in {"", "auto"}
provider_matches = _normalize_aux_provider(
_fallback_provider_from_label(str(actual_provider or ""))
) == _normalize_aux_provider(provider)
model_matches = str(actual_model or "").strip().lower() == model.lower()
certified = explicit_route and provider_matches and model_matches and effort == "none"
certified = explicit_route and provider_matches and model_matches and non_reasoning
if not certified:
return CompressionFastLane(False, None, None)
raw_cap = config.get("max_output_tokens")
try:
cap = 0 if isinstance(raw_cap, bool) else int(raw_cap)
except (TypeError, ValueError):
cap = 0
return CompressionFastLane(
True,
cap if cap > 0 else None,
cap,
{"enabled": False, "effort": "none"},
)
def _compression_config_claims_fast_lane(config: Dict[str, Any]) -> bool:
"""Whether task config declares fast-only controls that cannot leak."""
provider = str(config.get("provider") or "").strip().lower()
model = str(config.get("model") or "").strip().lower()
effort = str(config.get("reasoning_effort") or "").strip().lower()
raw_cap = config.get("max_output_tokens")
try:
cap = 0 if isinstance(raw_cap, bool) else int(raw_cap or 0)
except (TypeError, ValueError):
cap = 0
provider, model, non_reasoning, cap = _fast_lane_config_fields(config)
return (
provider not in {"", "auto"}
and model not in {"", "auto"}
and effort == "none"
and cap > 0
and model.lower() not in {"", "auto"}
and non_reasoning
and cap is not None
)
@@ -9854,11 +9890,15 @@ def _call_llm_impl(
tools=tools, timeout=effective_timeout, extra_body=effective_extra_body,
reasoning_config=reasoning_config,
base_url=_base_info or resolved_base_url, task=task)
if fast_compression_cap is not None:
if fast_compression_cap is not None and max_tokens is None:
# Normal auxiliary calls intentionally omit a cap on most
# OpenAI-compatible/local providers. This is the narrow exception:
# the configured compression route is concrete and certified
# non-reasoning, so a bounded summary request is intentional.
# ``max_tokens is None`` restricts the forced param to caps the
# certified lane itself produced — an explicit caller max_tokens is
# passed through untouched and keeps _build_call_kwargs's
# provider-quirk handling (same guard as the fallback path).
kwargs.update(auxiliary_max_tokens_param(fast_compression_cap, model=final_model))
if extra_headers:
kwargs["extra_headers"] = dict(extra_headers)
+97
View File
@@ -378,3 +378,100 @@ def test_fallback_cap_requires_independent_route_certification():
"enabled": False,
"effort": "none",
}
def test_reasoning_effort_aliases_certify_like_none():
"""Every spelling parse_reasoning_effort treats as disabled must certify.
_get_task_extra_body uses parse_reasoning_effort to disable reasoning for
"false"/"disabled"/YAML False exactly like "none"; the certification
predicate must agree or those users silently lose the fast lane.
"""
base = {"provider": "ollama", "model": "qwen3:8b", "max_output_tokens": 1400}
for alias in ("none", "false", "disabled", False):
lane = _resolve({**base, "reasoning_effort": alias})
assert lane.certified_non_reasoning is True, alias
assert lane.max_tokens == 1400, alias
# Empty/unset (provider default) and real efforts must NOT certify.
for not_disabled in ("", None, "low", "high", True):
lane = _resolve({**base, "reasoning_effort": not_disabled})
assert lane.certified_non_reasoning is False, not_disabled
assert lane.max_tokens is None, not_disabled
def test_timing_hooks_propagate_to_protected_call_worker_thread():
"""The protected daemon path must carry the timing hooks across threads.
_run_protected_sync_provider_call runs the provider callback on a daemon
worker. The dispatch/provider-response hooks are threading.local, so
without explicit propagation provider_dispatch_ms and
time_to_first_progress_ms silently vanish whenever compression takes the
protected path (the common case: aux_interrupt_protection + hard-cancel
source both active).
"""
from agent.auxiliary_client import (
_aux_timing_hook,
_aux_dispatch,
_aux_provider_response,
_notify_aux_dispatch,
_notify_aux_provider_response,
_run_protected_sync_provider_call,
aux_interrupt_protection,
)
seen = []
def _callback(_kwargs):
# Runs on the daemon worker thread — both notifies must reach the
# hooks installed on the owner thread.
_notify_aux_dispatch()
_notify_aux_provider_response()
return "ok"
with (
_aux_timing_hook(_aux_dispatch, lambda: seen.append("dispatch")),
_aux_timing_hook(_aux_provider_response, lambda: seen.append("response")),
aux_interrupt_protection(cancel_check=lambda: False),
):
result = _run_protected_sync_provider_call(_callback, {})
assert result == "ok"
assert "dispatch" in seen
assert "response" in seen
def test_explicit_caller_max_tokens_keeps_provider_quirk_handling():
"""An explicit caller cap must NOT be force-injected as a wire param.
_build_call_kwargs deliberately omits max_tokens for most
OpenAI-compatible providers (ZAI vision 400s on it; GPT-5/Copilot need
max_completion_tokens). Only a cap the certified lane itself produced may
bypass that handling. Before this guard, a caller-passed max_tokens on
the compression task flowed through _compression_fast_lane_controls as a
passthrough and was misread as a lane cap — forcing the param onto
providers where the omission was intentional (pre-fast-lane behavior).
"""
from agent.auxiliary_client import call_llm
config = {"provider": "auto", "model": "", "max_output_tokens": 0}
client = MagicMock()
client.base_url = "http://127.0.0.1:11434/v1"
response = object()
client.chat.completions.create.return_value = response
with (
patch("agent.auxiliary_client._get_auxiliary_task_config", return_value=config),
patch("agent.auxiliary_client._get_cached_client", return_value=(client, "qwen3:8b")),
patch("agent.auxiliary_client._validate_llm_response", return_value=response),
):
assert call_llm(
task="compression",
messages=[{"role": "user", "content": "summary request"}],
max_tokens=1500,
) is response
request = client.chat.completions.create.call_args.kwargs
assert "max_tokens" not in request
assert "max_completion_tokens" not in request