diff --git a/routes/chat_helpers.py b/routes/chat_helpers.py index 96ac103d8..eb5fcf489 100644 --- a/routes/chat_helpers.py +++ b/routes/chat_helpers.py @@ -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) diff --git a/routes/chat_routes.py b/routes/chat_routes.py index 09596ccbc..dccfaaba8 100644 --- a/routes/chat_routes.py +++ b/routes/chat_routes.py @@ -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: diff --git a/src/agent_loop.py b/src/agent_loop.py index 5de35898a..9b2e5e6d3 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -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 diff --git a/src/agent_runtime/context_resolution.py b/src/agent_runtime/context_resolution.py index e5064b518..8290f4837 100644 --- a/src/agent_runtime/context_resolution.py +++ b/src/agent_runtime/context_resolution.py @@ -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", diff --git a/src/clean_agent_preview.py b/src/clean_agent_preview.py index 78dce14ff..01b820bc5 100644 --- a/src/clean_agent_preview.py +++ b/src/clean_agent_preview.py @@ -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, diff --git a/src/context_compactor.py b/src/context_compactor.py index 1f9be5018..11f450323 100644 --- a/src/context_compactor.py +++ b/src/context_compactor.py @@ -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() diff --git a/tests/conftest.py b/tests/conftest.py index c73185eaa..cdc3129d9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/test_context_resolution.py b/tests/test_context_resolution.py index 4290c8aca..dcdde7bb0 100644 --- a/tests/test_context_resolution.py +++ b/tests/test_context_resolution.py @@ -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 diff --git a/tests/test_context_resolution_route.py b/tests/test_context_resolution_route.py new file mode 100644 index 000000000..9b64d04f9 --- /dev/null +++ b/tests/test_context_resolution_route.py @@ -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