mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-10-06 06:52:20 +02:00
fix(runtime): resolve compact-turn context once at the chat route
The first checkpoint removed terminal-metrics discovery, but a normal compact chat turn still ran two context systems: build_chat_context's legacy untyped lookup (directly or inside maybe_compact) and the typed resolver inside stream_preview. Resolve the typed ContextResolution once, at the chat route, before build_chat_context, using the session's provider credentials. The predicate mirrors _clean_v3_preview; every input it needs is known at that point and the native-workspace term cannot veto a requested clean route. The same object then: - sizes legacy history shaping in build_chat_context through a new maybe_compact(context_length=...) override, so no legacy probe runs; an unknown window still shapes with DEFAULT_CONTEXT but gains no provenance; - crosses stream_agent_loop (one new parameter, forwarded only at the compact dispatch) into stream_preview, which reuses it and probes only for callers that arrive without one or with one bound to another route. ContextResolution now records the endpoint and model it describes (endpoint URL excluded from repr and metrics). The bare legacy context_length is never converted into typed evidence. Credential scoping: origins compare with default ports normalized, an empty host is never trusted, and the probe client never follows redirects. Tests cover the configured origin, the server-resolved Tailscale form, scheme/port/lookalike/userinfo/path origins, redirects, and secret-free errors, logs and metrics. The conftest guard now replaces only the resolver's I/O edges (HTTP client and DNS-capable URL building) instead of the whole probe, and exposes a context_probe_ledger fixture, so route integration tests run the real resolver offline and can count metadata requests.
This commit is contained in:
Alexandre Teixeira
+13
-1
@@ -1027,11 +1027,15 @@ async def build_chat_context(
|
||||
persist_user_message: bool = True,
|
||||
interaction_mode: str | None = None,
|
||||
auto_escalated: bool = False,
|
||||
context_resolution=None,
|
||||
) -> ChatContext:
|
||||
"""Build the full context (preface + messages) for an LLM call.
|
||||
|
||||
This is the shared logic between /chat and /chat_stream — preset extraction,
|
||||
message preprocessing, memory/RAG/web injection, compaction, normalization.
|
||||
|
||||
``context_resolution`` is the turn's already resolved context window. When
|
||||
supplied, history shaping sizes against it instead of probing the endpoint.
|
||||
"""
|
||||
# Preset
|
||||
preset = extract_preset(chat_handler, preset_id)
|
||||
@@ -1225,12 +1229,20 @@ async def build_chat_context(
|
||||
# for every candidate. Running selected-model compaction here would mutate
|
||||
# session history before we know which route can answer and would make a
|
||||
# later larger-context candidate unable to recover discarded history.
|
||||
if context_resolution is not None:
|
||||
prepared_window = {"context_length": context_resolution.shaping_window}
|
||||
else:
|
||||
prepared_window = {}
|
||||
if defer_context_shaping:
|
||||
context_length = get_context_length(sess.endpoint_url, sess.model)
|
||||
context_length = (
|
||||
prepared_window.get("context_length")
|
||||
or get_context_length(sess.endpoint_url, sess.model)
|
||||
)
|
||||
was_compacted = False
|
||||
else:
|
||||
messages, context_length, was_compacted = await maybe_compact(
|
||||
sess, sess.endpoint_url, sess.model, messages, sess.headers, owner=user,
|
||||
**prepared_window,
|
||||
)
|
||||
_before_trim_messages = len(messages)
|
||||
_before_trim_tokens = estimate_tokens(messages)
|
||||
|
||||
@@ -2893,6 +2893,26 @@ def setup_chat_routes(
|
||||
allowed_models=_allowed_models_for_request(request),
|
||||
)
|
||||
|
||||
# A compact (clean v3) turn resolves its typed context window once,
|
||||
# here, with the session's provider credentials. History shaping below
|
||||
# and the compact runtime both reuse this exact object, so the turn
|
||||
# neither probes twice nor mixes the legacy untyped lookup into it.
|
||||
# The predicate mirrors ``_clean_v3_preview`` below; the native
|
||||
# workspace term cannot veto a requested clean route.
|
||||
_compact_context_resolution = None
|
||||
if _clean_v3_route_requested and _turn_contract_enabled(
|
||||
exact_tool_approval=exact_tool_approval,
|
||||
runtime_surface=str((client_runtime_context or {}).get("surface") or ""),
|
||||
native_workspace_contract=False,
|
||||
clean_v3_route=True,
|
||||
full_schema_route=(_effective_tool_schema_mode == "full"),
|
||||
):
|
||||
from src.agent_runtime.context_resolution import resolve_effective_context
|
||||
_compact_context_resolution = await resolve_effective_context(
|
||||
sess.endpoint_url, sess.model, headers=sess.headers,
|
||||
client_runtime_context=client_runtime_context,
|
||||
)
|
||||
|
||||
# Build shared context (stream path uses enhanced_message for context preface)
|
||||
ctx = await build_chat_context(
|
||||
sess, request, chat_handler, chat_processor,
|
||||
@@ -2923,6 +2943,7 @@ def setup_chat_routes(
|
||||
else None
|
||||
),
|
||||
persist_user_message=not tool_approval_continuation and not is_internal_tool_request(request),
|
||||
context_resolution=_compact_context_resolution,
|
||||
interaction_mode=chat_mode,
|
||||
auto_escalated=auto_escalated,
|
||||
)
|
||||
@@ -4459,6 +4480,9 @@ def setup_chat_routes(
|
||||
client_runtime_context=client_runtime_context,
|
||||
thinking_mode=thinking_mode,
|
||||
reasoning_effort=reasoning_effort,
|
||||
context_resolution=(
|
||||
_compact_context_resolution if _clean_v3_preview else None
|
||||
),
|
||||
):
|
||||
if chunk.startswith("data: ") and not chunk.startswith("data: [DONE]"):
|
||||
try:
|
||||
|
||||
@@ -20425,6 +20425,7 @@ async def stream_agent_loop(
|
||||
reasoning_effort: Optional[str] = None,
|
||||
_parent_run_id: Optional[str] = None,
|
||||
request_authority=MISSING_AUTHORITY,
|
||||
context_resolution=None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Streaming agent loop generator.
|
||||
|
||||
@@ -20477,6 +20478,9 @@ async def stream_agent_loop(
|
||||
max_rounds=max_rounds,
|
||||
max_tool_calls=max_tool_calls,
|
||||
temperature=temperature,
|
||||
# The route's prepared, typed context window; the compact runtime
|
||||
# must not resolve it a second time.
|
||||
context_resolution=context_resolution,
|
||||
):
|
||||
yield chunk
|
||||
return
|
||||
|
||||
@@ -24,7 +24,7 @@ Context sizing is not authority: nothing here grants or denies an operation.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass, replace
|
||||
from dataclasses import dataclass, field, replace
|
||||
from enum import Enum
|
||||
import hashlib
|
||||
import json
|
||||
@@ -90,6 +90,11 @@ class ContextResolution:
|
||||
provider_io: bool = False
|
||||
cached: bool = False
|
||||
probe_errors: tuple[str, ...] = ()
|
||||
# The route this resolution describes. Empty for resolutions built
|
||||
# directly from observations by internal callers. The URL can carry
|
||||
# credentials, so it stays out of repr() and to_dict().
|
||||
endpoint_url: str = field(default="", repr=False)
|
||||
model: str = ""
|
||||
|
||||
@property
|
||||
def mismatch(self) -> bool:
|
||||
@@ -100,6 +105,25 @@ class ContextResolution:
|
||||
"""Window the runtime may budget against; 0 means budget reactively."""
|
||||
return self.effective if self.evidence is not ContextEvidence.UNKNOWN else 0
|
||||
|
||||
@property
|
||||
def shaping_window(self) -> int:
|
||||
"""Window for the legacy history compaction/trim helpers.
|
||||
|
||||
Those helpers predate typed evidence and always size against some
|
||||
window, using DEFAULT_CONTEXT when none is known. This only feeds them
|
||||
a number; it never creates provenance for that number.
|
||||
"""
|
||||
if self.budget_limit:
|
||||
return self.budget_limit
|
||||
from src.model_context import DEFAULT_CONTEXT
|
||||
return DEFAULT_CONTEXT
|
||||
|
||||
def applies_to(self, endpoint_url: str, model: str) -> bool:
|
||||
"""Whether this resolution may be reused for the given route."""
|
||||
if not self.endpoint_url and not self.model:
|
||||
return True
|
||||
return self.endpoint_url == endpoint_url and self.model == model
|
||||
|
||||
def observe_runtime_limit(self, limit: Any, source: str = "provider_rejection") -> "ContextResolution":
|
||||
"""Fold a limit the provider stated during this turn. Performs no I/O."""
|
||||
try:
|
||||
@@ -117,6 +141,8 @@ class ContextResolution:
|
||||
provider_io=self.provider_io,
|
||||
cached=self.cached,
|
||||
probe_errors=self.probe_errors,
|
||||
endpoint_url=self.endpoint_url,
|
||||
model=self.model,
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
@@ -209,14 +235,40 @@ def clear_probe_cache() -> None:
|
||||
_probe_cache.clear()
|
||||
|
||||
|
||||
_DEFAULT_PORTS = {"http": 80, "https": 443}
|
||||
|
||||
|
||||
def _origin(url: str) -> tuple[str, str, Optional[int]]:
|
||||
parsed = urlparse(url or "")
|
||||
return (parsed.scheme.lower(), (parsed.hostname or "").lower(), parsed.port)
|
||||
scheme = parsed.scheme.lower()
|
||||
try:
|
||||
port = parsed.port
|
||||
except ValueError:
|
||||
return ("", "", None)
|
||||
return (scheme, (parsed.hostname or "").lower(), port or _DEFAULT_PORTS.get(scheme))
|
||||
|
||||
|
||||
def _http_client(timeout: float):
|
||||
# Credentials must never follow a redirect to another location.
|
||||
return httpx.AsyncClient(timeout=timeout, follow_redirects=False)
|
||||
|
||||
|
||||
def _provider_urls(endpoint_url: str) -> tuple[Optional[str], str]:
|
||||
"""Models catalog URL and the server-resolved form of the endpoint.
|
||||
|
||||
Both come from the existing endpoint resolver, which may rewrite an
|
||||
unresolvable host to its Tailscale address. Blocking (DNS, subprocess);
|
||||
call it off the event loop.
|
||||
"""
|
||||
from src.endpoint_resolver import build_models_url, resolve_url
|
||||
|
||||
return build_models_url(endpoint_url), resolve_url(endpoint_url)
|
||||
|
||||
|
||||
def _probe_headers(trusted_origins, target_url: str, headers: Optional[Mapping[str, Any]]) -> dict:
|
||||
"""Forward the turn's provider credentials only to the provider's origin."""
|
||||
if not headers or _origin(target_url) not in trusted_origins:
|
||||
origin = _origin(target_url)
|
||||
if not headers or not origin[1] or origin not in trusted_origins:
|
||||
return {}
|
||||
return {
|
||||
str(name): str(value) for name, value in headers.items()
|
||||
@@ -268,11 +320,12 @@ def _positive_int(value) -> int:
|
||||
|
||||
async def _probe(endpoint_url, model, headers, is_local, observations, errors, timeout):
|
||||
from src.copilot import is_copilot_base
|
||||
from src.endpoint_resolver import build_models_url, resolve_url
|
||||
from src.model_context import _model_ctx_from_entry
|
||||
|
||||
# Credentials go only to the configured provider's origin, or to the
|
||||
# form of that same endpoint the server-owned resolver produced.
|
||||
trusted = {_origin(endpoint_url)}
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
async with _http_client(timeout) as client:
|
||||
if is_local:
|
||||
base = _serving_base(endpoint_url)
|
||||
slots = await _get_json(
|
||||
@@ -303,13 +356,12 @@ async def _probe(endpoint_url, model, headers, is_local, observations, errors, t
|
||||
errors.append("models:unsupported_endpoint")
|
||||
return
|
||||
# URL building may resolve the host (DNS, tailscale lookup); keep that
|
||||
# off the event loop and inside the probe deadline. The resolved host
|
||||
# is the same provider the chat request reaches.
|
||||
models_url = await asyncio.to_thread(build_models_url, endpoint_url)
|
||||
# off the event loop and inside the probe deadline.
|
||||
models_url, resolved_endpoint = await asyncio.to_thread(_provider_urls, endpoint_url)
|
||||
if not models_url:
|
||||
errors.append("models:unsupported_endpoint")
|
||||
return
|
||||
trusted.add(_origin(await asyncio.to_thread(resolve_url, endpoint_url)))
|
||||
trusted.add(_origin(resolved_endpoint))
|
||||
payload = await _get_json(
|
||||
client, models_url, _probe_headers(trusted, models_url, headers),
|
||||
errors, "models",
|
||||
@@ -440,7 +492,10 @@ async def resolve_effective_context(
|
||||
if known:
|
||||
observations.append(ContextObservation(ContextEvidence.KNOWN_TABLE, int(known), "known_table"))
|
||||
resolution = combine_observations(observations)
|
||||
resolution = replace(resolution, provider_io=provider_io, cached=cached, probe_errors=errors)
|
||||
resolution = replace(
|
||||
resolution, provider_io=provider_io, cached=cached, probe_errors=errors,
|
||||
endpoint_url=endpoint_url or "", model=model or "",
|
||||
)
|
||||
if resolution.mismatch:
|
||||
logger.info(
|
||||
"Context window sources disagree for %s: %s",
|
||||
|
||||
@@ -5254,8 +5254,10 @@ async def stream_preview(*, endpoint_url, model, messages, headers, turn_contrac
|
||||
calendar_create_confirmation = ''
|
||||
# Resolve the effective window once, before any model request, so the
|
||||
# turn budgets against it and metrics report exactly what it ran under.
|
||||
# Terminal metrics must never start this discovery themselves.
|
||||
if context_resolution is None:
|
||||
# The chat route prepares it; only callers arriving without one (or with
|
||||
# one for a different route) resolve here. Terminal metrics must never
|
||||
# start this discovery themselves.
|
||||
if context_resolution is None or not context_resolution.applies_to(endpoint_url, model):
|
||||
from src.agent_runtime.context_resolution import resolve_effective_context
|
||||
context_resolution = await resolve_effective_context(
|
||||
endpoint_url, model, headers=headers,
|
||||
|
||||
@@ -481,12 +481,17 @@ async def maybe_compact(
|
||||
persist: bool = True,
|
||||
compaction_state: Optional[Dict[str, Any]] = None,
|
||||
deterministic: bool = False,
|
||||
context_length: Optional[int] = None,
|
||||
) -> tuple:
|
||||
"""Check context usage and compact if above threshold.
|
||||
|
||||
``context_length`` lets a caller that already resolved the turn's window
|
||||
supply it, so this helper does not query the endpoint a second time.
|
||||
|
||||
Returns (messages, context_length, was_compacted).
|
||||
"""
|
||||
context_length = get_context_length(endpoint_url, model)
|
||||
if context_length is None:
|
||||
context_length = get_context_length(endpoint_url, model)
|
||||
used = estimate_tokens(messages)
|
||||
pct = (used / context_length) * 100 if context_length else 0
|
||||
threshold = auto_compact_threshold_percent()
|
||||
|
||||
+40
-11
@@ -205,32 +205,61 @@ def _no_leaked_module_stubs():
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_context_window_network_probe(request):
|
||||
"""Keep turn preparation from probing real provider metadata in tests.
|
||||
"""Keep the turn context-window resolver offline in tests.
|
||||
|
||||
The compact runtime resolves its context window before the first model
|
||||
request. Tests that drive it with placeholder endpoints must not perform
|
||||
DNS or HTTP lookups; modules that exercise the probe opt in with a
|
||||
module-level ``CONTEXT_PROBE_NETWORK = True`` and supply their own client.
|
||||
Compact turns resolve their window before the first model request, and
|
||||
most tests drive them with placeholder endpoints. Only the resolver's two
|
||||
I/O edges are replaced: URL resolution (DNS/Tailscale lookups) and its
|
||||
HTTP client, which records each attempted metadata request and fails it
|
||||
as a transport error. Everything else (route wiring, caching, credential
|
||||
scoping, evidence selection) runs for real, so an unintended extra probe
|
||||
stays visible through the ``context_probe_ledger`` fixture.
|
||||
|
||||
Modules that install their own fake client opt out with a module-level
|
||||
``CONTEXT_PROBE_NETWORK = True``.
|
||||
"""
|
||||
ledger = []
|
||||
if getattr(request.module, "CONTEXT_PROBE_NETWORK", False):
|
||||
yield
|
||||
yield ledger
|
||||
return
|
||||
try:
|
||||
from src.agent_runtime import context_resolution
|
||||
except Exception:
|
||||
yield
|
||||
yield ledger
|
||||
return
|
||||
|
||||
async def _disabled_probe(endpoint_url, model, headers, is_local, observations, errors, timeout):
|
||||
errors.append("probe_disabled_in_tests")
|
||||
class _OfflineMetadataClient:
|
||||
def __init__(self, timeout=None):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
|
||||
async def get(self, url, headers=None):
|
||||
ledger.append({"url": url, "headers": dict(headers or {})})
|
||||
raise context_resolution.httpx.ConnectError("network disabled in tests")
|
||||
|
||||
def _offline_provider_urls(endpoint_url):
|
||||
base = endpoint_url.split("/v1")[0] if "/v1" in endpoint_url else endpoint_url.rstrip("/")
|
||||
return base + "/v1/models", endpoint_url
|
||||
|
||||
# A private patcher keeps the shared ``monkeypatch`` fixture's teardown
|
||||
# order unchanged for tests that check their own sys.modules hygiene.
|
||||
patcher = pytest.MonkeyPatch()
|
||||
patcher.setattr(context_resolution, "_probe", _disabled_probe)
|
||||
patcher.setattr(context_resolution, "_http_client", _OfflineMetadataClient)
|
||||
patcher.setattr(context_resolution, "_provider_urls", _offline_provider_urls)
|
||||
context_resolution.clear_probe_cache()
|
||||
try:
|
||||
yield
|
||||
yield ledger
|
||||
finally:
|
||||
patcher.undo()
|
||||
context_resolution.clear_probe_cache()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def context_probe_ledger(_no_context_window_network_probe):
|
||||
"""Metadata requests the context resolver attempted during this test."""
|
||||
return _no_context_window_network_probe
|
||||
|
||||
@@ -89,13 +89,14 @@ class _FakeNetwork:
|
||||
self.get_error = get_error
|
||||
self.calls = []
|
||||
self.requests = []
|
||||
self.client_kwargs = []
|
||||
|
||||
def client_factory(self):
|
||||
network = self
|
||||
|
||||
class Client:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
network.client_kwargs.append(kwargs)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
@@ -511,3 +512,110 @@ async def test_compact_turn_proceeds_when_metadata_probe_times_out(monkeypatch):
|
||||
assert resolution["probe_errors"] == ["models:timeout"]
|
||||
assert metrics["context_length"] == 65536
|
||||
assert resolution["evidence"] == "operator_declared"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Credential scoping (adversarial)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _models_at(monkeypatch, models_url, resolved=None):
|
||||
monkeypatch.setattr("src.endpoint_resolver.build_models_url", lambda base: models_url)
|
||||
monkeypatch.setattr(
|
||||
"src.endpoint_resolver.resolve_url", lambda url: url if resolved is None else resolved,
|
||||
)
|
||||
|
||||
|
||||
async def _forwarded_headers(monkeypatch, endpoint_url, models_url, resolved=None):
|
||||
network = _install(monkeypatch, _FakeNetwork({"/models": _catalog("acme-model", max_model_len=8192)}))
|
||||
_models_at(monkeypatch, models_url, resolved)
|
||||
await resolve_effective_context(endpoint_url, "acme-model", headers=AUTH)
|
||||
[(kind, url, headers)] = network.calls
|
||||
assert url == models_url
|
||||
return headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("models_url", [
|
||||
"https://provider.test/v1/models",
|
||||
"https://PROVIDER.test:443/v1/models",
|
||||
])
|
||||
async def test_credentials_reach_the_configured_provider_origin(monkeypatch, models_url):
|
||||
headers = await _forwarded_headers(
|
||||
monkeypatch, "https://provider.test/v1/chat/completions", models_url,
|
||||
)
|
||||
assert headers == {"Authorization": "Bearer secret-token"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credentials_reach_only_the_server_resolved_form_of_the_provider(monkeypatch):
|
||||
endpoint = "http://gpu-box:8000/v1/chat/completions"
|
||||
resolved = "http://100.64.0.9:8000/v1/chat/completions"
|
||||
assert await _forwarded_headers(
|
||||
monkeypatch, endpoint, "http://100.64.0.9:8000/v1/models", resolved,
|
||||
) == {"Authorization": "Bearer secret-token"}
|
||||
cr.clear_probe_cache()
|
||||
# The same address is not trusted when the server resolver did not
|
||||
# produce it for this endpoint.
|
||||
assert await _forwarded_headers(
|
||||
monkeypatch, endpoint, "http://100.64.0.9:8000/v1/models",
|
||||
) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("models_url", [
|
||||
"http://provider.test/v1/models", # scheme downgrade
|
||||
"https://provider.test:8443/v1/models", # other port
|
||||
"https://provider.test.evil.example/v1/models", # lookalike host
|
||||
"https://evil.example/provider.test/v1/models", # host in path
|
||||
"https://provider.test@evil.example/v1/models", # host in userinfo
|
||||
"/v1/models", # no origin at all
|
||||
])
|
||||
async def test_unrelated_models_url_receives_no_credentials(monkeypatch, models_url):
|
||||
assert await _forwarded_headers(
|
||||
monkeypatch, "https://provider.test/v1/chat/completions", models_url,
|
||||
) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_client_never_follows_redirects(monkeypatch):
|
||||
network = _install(monkeypatch, _FakeNetwork({"/models": _Response(302, None)}))
|
||||
resolution = await resolve_effective_context(REMOTE, "acme-model", headers=AUTH)
|
||||
assert network.client_kwargs and all(
|
||||
kwargs.get("follow_redirects") is False for kwargs in network.client_kwargs
|
||||
)
|
||||
assert resolution.probe_errors == ("models:http_302",)
|
||||
assert len(network.calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", [
|
||||
_FakeNetwork(get_error=httpx.ConnectError(
|
||||
"connect to https://user:pw-secret@provider.test/v1/models?api_key=query-secret failed")),
|
||||
_FakeNetwork(get_error=httpx.ReadTimeout("Bearer secret-token timed out")),
|
||||
_FakeNetwork({"/models": _Response(401, None)}),
|
||||
_FakeNetwork({"/models": _Response(200, ValueError("api_key=query-secret"))}),
|
||||
_FakeNetwork(get_error=RuntimeError("Authorization: Bearer secret-token")),
|
||||
])
|
||||
async def test_probe_errors_never_expose_credentials_or_urls(monkeypatch, caplog, failure):
|
||||
import logging
|
||||
|
||||
caplog.set_level(logging.DEBUG)
|
||||
endpoint = "https://user:pw-secret@provider.test/v1/chat/completions?api_key=query-secret"
|
||||
_install(monkeypatch, failure)
|
||||
resolution = await resolve_effective_context(endpoint, "acme-model", headers=AUTH)
|
||||
assert resolution.probe_errors
|
||||
exposed = json.dumps(resolution.to_dict()) + caplog.text + json.dumps(
|
||||
context_metrics(resolution, 10),
|
||||
)
|
||||
for secret in ("secret-token", "pw-secret", "query-secret", "provider.test", "Authorization"):
|
||||
assert secret not in exposed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bound_endpoint_url_stays_out_of_repr_and_metrics(monkeypatch):
|
||||
_install(monkeypatch, _FakeNetwork({"/models": _Response(503, None)}))
|
||||
endpoint = "https://user:pw-secret@provider.test/v1/chat/completions?api_key=query-secret"
|
||||
resolution = await resolve_effective_context(endpoint, "acme-model", headers=AUTH)
|
||||
assert resolution.applies_to(endpoint, "acme-model")
|
||||
for rendered in (repr(resolution), str(resolution), json.dumps(resolution.to_dict())):
|
||||
assert "pw-secret" not in rendered and "query-secret" not in rendered
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
"""One typed context resolution per compact chat turn, end to end.
|
||||
|
||||
These tests keep the conftest offline guard active on purpose: the real route,
|
||||
agent loop and compact runtime run, only the resolver's HTTP/DNS edges are
|
||||
offline, and ``context_probe_ledger`` records every metadata request.
|
||||
"""
|
||||
from dataclasses import replace
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import src.model_context as model_context
|
||||
from src.agent_runtime import context_resolution as cr
|
||||
from src.agent_runtime.context_resolution import (
|
||||
ContextEvidence,
|
||||
ContextObservation,
|
||||
combine_observations,
|
||||
)
|
||||
from tests.test_foreground_model_routing import _RouteRequest, _chat_stream_endpoint
|
||||
|
||||
COMPACT_MODEL = "odysseus-qwen3.5-tools-pre-heretic"
|
||||
|
||||
|
||||
class _ModelStream:
|
||||
def __init__(self, lines, status=200, text=""):
|
||||
self.status_code = status
|
||||
self.text = text
|
||||
self._lines = lines
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
|
||||
async def aread(self):
|
||||
return self.text.encode()
|
||||
|
||||
def raise_for_status(self):
|
||||
if self.status_code >= 400:
|
||||
raise AssertionError(f"unexpected provider status {self.status_code}")
|
||||
|
||||
async def aiter_lines(self):
|
||||
for line in self._lines:
|
||||
yield line
|
||||
|
||||
|
||||
def _answer(prompt_tokens=1024):
|
||||
return _ModelStream([
|
||||
"data: " + json.dumps({"choices": [{"delta": {"content": "Hello."}}]}),
|
||||
"data: " + json.dumps({"choices": [{"delta": {}}],
|
||||
"usage": {"prompt_tokens": prompt_tokens, "completion_tokens": 2}}),
|
||||
"data: [DONE]",
|
||||
])
|
||||
|
||||
|
||||
def _install_model(monkeypatch, *responses):
|
||||
"""Fake provider for chat completions only; metadata goes to the ledger."""
|
||||
import src.clean_agent_preview as preview
|
||||
|
||||
queue = list(responses)
|
||||
sent = []
|
||||
|
||||
class Client:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
|
||||
def stream(self, method, url, headers=None, json=None):
|
||||
sent.append(json)
|
||||
return queue.pop(0) if queue else _answer()
|
||||
|
||||
monkeypatch.setattr(preview.httpx, "AsyncClient", Client)
|
||||
return sent
|
||||
|
||||
|
||||
def _forbid_legacy_probe(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def legacy(*args, **kwargs):
|
||||
calls.append(args)
|
||||
raise AssertionError("legacy context probe used on a compact turn")
|
||||
|
||||
monkeypatch.setattr(model_context, "_query_context_length", legacy)
|
||||
model_context._context_cache.clear()
|
||||
return calls
|
||||
|
||||
|
||||
def _spy(monkeypatch):
|
||||
"""Count resolutions and capture what the compact runtime received/emitted."""
|
||||
import src.clean_agent_preview as preview
|
||||
|
||||
seen = {"resolutions": [], "preview_kwargs": [], "preview_chunks": []}
|
||||
real_resolve = cr.resolve_effective_context
|
||||
|
||||
async def counting_resolve(*args, **kwargs):
|
||||
result = await real_resolve(*args, **kwargs)
|
||||
seen["resolutions"].append(result)
|
||||
return result
|
||||
|
||||
real_preview = preview.stream_preview
|
||||
|
||||
async def recording_preview(**kwargs):
|
||||
seen["preview_kwargs"].append(kwargs)
|
||||
async for chunk in real_preview(**kwargs):
|
||||
seen["preview_chunks"].append(chunk)
|
||||
yield chunk
|
||||
|
||||
monkeypatch.setattr(cr, "resolve_effective_context", counting_resolve)
|
||||
monkeypatch.setattr(preview, "stream_preview", recording_preview)
|
||||
return seen
|
||||
|
||||
|
||||
def _metrics(chunks):
|
||||
for chunk in chunks:
|
||||
if chunk.startswith("data: {"):
|
||||
event = json.loads(chunk[6:])
|
||||
if event.get("type") == "metrics":
|
||||
return event["data"]
|
||||
raise AssertionError("no metrics event")
|
||||
|
||||
|
||||
async def _drive_route(monkeypatch, *, model=COMPACT_MODEL, message="hello"):
|
||||
from routes import chat_routes
|
||||
import src.agent_loop as agent_loop
|
||||
|
||||
captured = {}
|
||||
endpoint = _chat_stream_endpoint(
|
||||
monkeypatch, "agent", captured, capture_context=True, session_model=model,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
chat_routes, "coerce_message_and_session", lambda *args, **kwargs: (message, "session-1"),
|
||||
)
|
||||
# The real agent loop, compact dispatch and compact runtime run below.
|
||||
monkeypatch.setattr(chat_routes, "stream_agent_loop", agent_loop.stream_agent_loop)
|
||||
request = _RouteRequest("agent")
|
||||
request._form.update({"message": message, "compare_mode": "false"})
|
||||
response = await endpoint(request)
|
||||
body = [chunk async for chunk in response.body_iterator]
|
||||
return captured, body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compact_chat_route_resolves_once_and_reuses_the_exact_object(
|
||||
monkeypatch, context_probe_ledger,
|
||||
):
|
||||
legacy = _forbid_legacy_probe(monkeypatch)
|
||||
seen = _spy(monkeypatch)
|
||||
sent = _install_model(monkeypatch)
|
||||
|
||||
captured, _ = await _drive_route(monkeypatch)
|
||||
|
||||
# One typed resolution, one metadata request, no legacy lookup.
|
||||
[resolution] = seen["resolutions"]
|
||||
assert len(context_probe_ledger) == 1
|
||||
assert context_probe_ledger[0]["url"] == "https://selected.example/v1/models"
|
||||
# The route used the session's provider credentials.
|
||||
assert context_probe_ledger[0]["headers"] == {"Authorization": "Bearer selected"}
|
||||
assert legacy == []
|
||||
assert sent, "the compact runtime never reached the model"
|
||||
|
||||
# The exact object crosses route -> build_chat_context and
|
||||
# route -> stream_agent_loop -> stream_preview.
|
||||
assert captured["build_context"]["context_resolution"] is resolution
|
||||
[preview_kwargs] = seen["preview_kwargs"]
|
||||
assert preview_kwargs["context_resolution"] is resolution
|
||||
assert resolution.applies_to("https://selected.example/v1", COMPACT_MODEL)
|
||||
|
||||
# Metrics report that same resolution; the offline probe failed, so the
|
||||
# trusted table supplies the window.
|
||||
metrics = _metrics(seen["preview_chunks"])
|
||||
assert metrics["context_resolution"] == resolution.to_dict()
|
||||
assert metrics["context_length"] == 131072
|
||||
assert metrics["context_resolution"]["evidence"] == "known_table"
|
||||
assert seen["preview_chunks"][-1] == "data: [DONE]\n\n"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compact_route_metrics_fold_a_provider_limit_without_a_second_probe(
|
||||
monkeypatch, context_probe_ledger,
|
||||
):
|
||||
_forbid_legacy_probe(monkeypatch)
|
||||
seen = _spy(monkeypatch)
|
||||
rejection = (
|
||||
"This model's maximum context length is 4096 tokens. However, you requested "
|
||||
"5000 tokens (4800 in the messages, 200 in the completion)."
|
||||
)
|
||||
_install_model(
|
||||
monkeypatch,
|
||||
_ModelStream([], status=400, text=json.dumps({"error": {"message": rejection}})),
|
||||
_answer(),
|
||||
)
|
||||
|
||||
await _drive_route(monkeypatch)
|
||||
|
||||
[resolution] = seen["resolutions"]
|
||||
assert len(context_probe_ledger) == 1
|
||||
metrics = _metrics(seen["preview_chunks"])
|
||||
reported = metrics["context_resolution"]
|
||||
assert metrics["context_length"] == 4096
|
||||
assert (reported["evidence"], reported["source"]) == ("runtime_confirmed", "provider_rejection")
|
||||
# Everything else is the prepared resolution, unchanged.
|
||||
assert reported == resolution.observe_runtime_limit(4096).to_dict()
|
||||
assert reported["observations"][:-1] == resolution.to_dict()["observations"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regular_model_route_does_not_prepare_a_compact_resolution(
|
||||
monkeypatch, context_probe_ledger,
|
||||
):
|
||||
from routes import chat_routes
|
||||
|
||||
seen = _spy(monkeypatch)
|
||||
captured = {}
|
||||
endpoint = _chat_stream_endpoint(monkeypatch, "agent", captured, capture_context=True)
|
||||
|
||||
async def capture_agent(*args, **kwargs):
|
||||
captured["agent_kwargs"] = kwargs
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
monkeypatch.setattr(chat_routes, "stream_agent_loop", capture_agent)
|
||||
request = _RouteRequest("agent")
|
||||
request._form.update({"message": "hello", "compare_mode": "false"})
|
||||
response = await endpoint(request)
|
||||
async for _ in response.body_iterator:
|
||||
pass
|
||||
|
||||
assert seen["resolutions"] == []
|
||||
assert context_probe_ledger == []
|
||||
assert captured["build_context"]["context_resolution"] is None
|
||||
assert captured["agent_kwargs"].get("context_resolution") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bare_context_length_never_becomes_typed_provenance(monkeypatch, context_probe_ledger):
|
||||
"""A legacy integer handed to the agent loop is not evidence of anything."""
|
||||
import src.agent_loop as agent_loop
|
||||
from src.tool_policy import ToolPolicy
|
||||
from src.turn_contract import resolve_full_inventory_contract
|
||||
from src.clean_agent_preview import MODE
|
||||
|
||||
seen = _spy(monkeypatch)
|
||||
_install_model(monkeypatch)
|
||||
contract = replace(
|
||||
resolve_full_inventory_contract(schemas=[], policy=ToolPolicy()), selection_mode=MODE,
|
||||
)
|
||||
chunks = [chunk async for chunk in agent_loop.stream_agent_loop(
|
||||
"https://selected.example/v1", "mystery-model", [{"role": "user", "content": "hello"}],
|
||||
headers={"Authorization": "Bearer selected"}, context_length=4096,
|
||||
turn_contract=contract, session_id="s", owner="alice",
|
||||
tool_policy=ToolPolicy(), disabled_tools=set(),
|
||||
)]
|
||||
|
||||
# A direct caller without a prepared resolution still resolves safely, once.
|
||||
[resolution] = seen["resolutions"]
|
||||
assert len(context_probe_ledger) == 1
|
||||
metrics = _metrics(chunks)
|
||||
assert metrics["context_length"] == 0
|
||||
assert metrics["context_resolution"]["evidence"] == "unknown"
|
||||
assert all(
|
||||
observation["value"] != 4096
|
||||
for observation in metrics["context_resolution"]["observations"]
|
||||
)
|
||||
assert resolution.to_dict() == metrics["context_resolution"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_preview_rejects_a_resolution_prepared_for_another_route(
|
||||
monkeypatch, context_probe_ledger,
|
||||
):
|
||||
import src.clean_agent_preview as preview
|
||||
from src.tool_policy import ToolPolicy
|
||||
from src.turn_contract import resolve_full_inventory_contract
|
||||
|
||||
_install_model(monkeypatch)
|
||||
foreign = combine_observations([
|
||||
ContextObservation(ContextEvidence.PROVIDER_ADVERTISED, 2048, "models_catalog"),
|
||||
])
|
||||
foreign = replace(foreign, endpoint_url="https://other.example/v1", model="other-model")
|
||||
chunks = [chunk async for chunk in preview.stream_preview(
|
||||
endpoint_url="https://selected.example/v1", model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hello"}], headers={},
|
||||
turn_contract=resolve_full_inventory_contract(schemas=[], policy=ToolPolicy()),
|
||||
session_id="s", owner="alice", disabled_tools=set(), tool_policy=ToolPolicy(),
|
||||
context_resolution=foreign,
|
||||
)]
|
||||
metrics = _metrics(chunks)
|
||||
assert len(context_probe_ledger) == 1
|
||||
assert metrics["context_length"] == 128000
|
||||
assert metrics["context_resolution"]["evidence"] == "known_table"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_supplied_resolution_means_no_probe_in_stream_preview(monkeypatch, context_probe_ledger):
|
||||
import src.clean_agent_preview as preview
|
||||
from src.tool_policy import ToolPolicy
|
||||
from src.turn_contract import resolve_full_inventory_contract
|
||||
|
||||
_install_model(monkeypatch)
|
||||
prepared = await cr.resolve_effective_context(
|
||||
"https://selected.example/v1", "gpt-4o", headers={"Authorization": "Bearer selected"},
|
||||
)
|
||||
assert len(context_probe_ledger) == 1
|
||||
chunks = [chunk async for chunk in preview.stream_preview(
|
||||
endpoint_url="https://selected.example/v1", model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hello"}], headers={"Authorization": "Bearer selected"},
|
||||
turn_contract=resolve_full_inventory_contract(schemas=[], policy=ToolPolicy()),
|
||||
session_id="s", owner="alice", disabled_tools=set(), tool_policy=ToolPolicy(),
|
||||
context_resolution=prepared,
|
||||
)]
|
||||
# Neither preparation nor terminal metrics probed again.
|
||||
assert len(context_probe_ledger) == 1
|
||||
assert _metrics(chunks)["context_resolution"] == prepared.to_dict()
|
||||
assert chunks[-1] == "data: [DONE]\n\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build_chat_context consumes the prepared resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _context_harness(monkeypatch):
|
||||
from tests.test_kv_cache_invalidation_2927 import _build_context_harness, _install_chat_helpers_stubs
|
||||
|
||||
chat_helpers = _install_chat_helpers_stubs(monkeypatch)
|
||||
sess, request, chat_handler, chat_processor = _build_context_harness(monkeypatch, chat_helpers, history=[])
|
||||
legacy = []
|
||||
monkeypatch.setattr(
|
||||
chat_helpers, "get_context_length",
|
||||
lambda *args: legacy.append(args) or 8192,
|
||||
)
|
||||
compactions = []
|
||||
|
||||
async def recording_maybe_compact(sess, endpoint_url, model, messages, headers, owner=None, **kwargs):
|
||||
compactions.append(kwargs)
|
||||
return messages, kwargs.get("context_length", 8192), False
|
||||
|
||||
monkeypatch.setattr(chat_helpers, "maybe_compact", recording_maybe_compact)
|
||||
return chat_helpers, (sess, request, chat_handler, chat_processor), legacy, compactions
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("defer", [False, True])
|
||||
async def test_build_chat_context_shapes_with_the_prepared_resolution(monkeypatch, defer):
|
||||
chat_helpers, (sess, request, handler, processor), legacy, compactions = _context_harness(monkeypatch)
|
||||
prepared = combine_observations([
|
||||
ContextObservation(ContextEvidence.PROVIDER_ADVERTISED, 32768, "models_catalog"),
|
||||
])
|
||||
ctx = await chat_helpers.build_chat_context(
|
||||
sess=sess, request=request, chat_handler=handler, chat_processor=processor,
|
||||
message="hello", session_id="s", defer_context_shaping=defer,
|
||||
context_resolution=prepared,
|
||||
)
|
||||
assert legacy == []
|
||||
assert ctx.context_length == 32768
|
||||
assert compactions == ([] if defer else [{"context_length": 32768}])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_chat_context_unknown_resolution_shapes_with_legacy_default(monkeypatch):
|
||||
chat_helpers, (sess, request, handler, processor), legacy, compactions = _context_harness(monkeypatch)
|
||||
ctx = await chat_helpers.build_chat_context(
|
||||
sess=sess, request=request, chat_handler=handler, chat_processor=processor,
|
||||
message="hello", session_id="s", context_resolution=combine_observations([]),
|
||||
)
|
||||
# Shaping still needs a number, but no probe and no provenance is created.
|
||||
assert legacy == []
|
||||
assert ctx.context_length == model_context.DEFAULT_CONTEXT
|
||||
assert compactions == [{"context_length": model_context.DEFAULT_CONTEXT}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_chat_context_without_resolution_keeps_legacy_lookup(monkeypatch):
|
||||
chat_helpers, (sess, request, handler, processor), legacy, compactions = _context_harness(monkeypatch)
|
||||
await chat_helpers.build_chat_context(
|
||||
sess=sess, request=request, chat_handler=handler, chat_processor=processor,
|
||||
message="hello", session_id="s", defer_context_shaping=True,
|
||||
)
|
||||
assert len(legacy) == 1
|
||||
assert compactions == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_offline_guard_replaces_only_io_edges(context_probe_ledger):
|
||||
"""The conftest guard must not mask the resolver itself."""
|
||||
headers = {"Authorization": "Bearer selected", "Content-Type": "application/json"}
|
||||
first = await cr.resolve_effective_context("https://selected.example/v1", "gpt-4o", headers=headers)
|
||||
second = await cr.resolve_effective_context("https://selected.example/v1", "gpt-4o", headers=headers)
|
||||
# Real request construction and credential scoping ran...
|
||||
assert context_probe_ledger == [{
|
||||
"url": "https://selected.example/v1/models",
|
||||
"headers": {"Authorization": "Bearer selected"},
|
||||
}]
|
||||
# ...as did real error mapping, evidence selection and caching.
|
||||
assert first.probe_errors == ("models:transport_error",)
|
||||
assert (first.evidence, first.effective) == (ContextEvidence.KNOWN_TABLE, 128000)
|
||||
assert second.cached and not second.provider_io
|
||||
Reference in New Issue
Block a user