From e86f3c5dc7329a75fcf1d158130818d61d247378 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 06:20:45 -0700 Subject: [PATCH] Studio: wire OpenAI Responses server-side context compaction (#5687) * Studio: wire OpenAI Responses server-side context compaction The OpenAI Responses API accepts a `context_management` field that enables server-side compaction. When the rendered prompt crosses the configured threshold, the API runs a server-side compaction step and the request continues against the compacted prefix. No beta header and no dated version pin are required, per the docs. Changes: - Add `compaction_threshold: Optional[int]` (ge=1_000, le=2_000_000) to ChatCompletionRequest. Thread through `routes/inference.py` -> `stream_chat_completion` -> `_stream_openai_responses`. - In `_stream_openai_responses`, when threshold is set AND the base URL points at cloud OpenAI (api.openai.com), attach `context_management: [{type:"compaction", compact_threshold:N}]` to the outbound body. Non-cloud bases (ollama, llama.cpp, "custom" presets) silently drop the field so we don't 400 those servers. - Add `test_openai_compaction.py` with 4 cases: cloud OpenAI sets the field verbatim, low-threshold probe passes through (we don't clamp on the OpenAI side because the API accepts whatever), non-cloud base drops the field, omitted threshold leaves body untouched. Live verified against the real OpenAI API on gpt-5.5: `context_management:[{type:"compaction", compact_threshold:200000}]` returns 200 with no error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: accept Azure OpenAI base URLs + raise compaction floor Two reviewer follow-ups on the OpenAI compaction PR: 1. The `is_openai_cloud = "api.openai.com" in self.base_url` check excluded Azure OpenAI Foundry, even though Azure exposes the same /v1/responses extensions (context_management, prompt_cache_retention, container shell). Users on Azure saw their compaction toggle silently no-op. Broadened the check to also match `*.openai.azure.com` and made it case-insensitive so URLs copy-pasted from the Azure portal still resolve. Non-cloud OpenAI-compatible servers (ollama / llama.cpp / vLLM / "custom" preset) still fall outside the gate. 2. The schema floor on compaction_threshold was ge=1_000, which is well below the upstream Responses API's effective minimum (vercel/ai#12486, langchain-ai/langchain#35464 report `compact_threshold is not enabled` 400s on Azure at 100k; cloud uses 200k as the canonical example). Raised the floor to 10k so obvious typos surface as a clean 422 from FastAPI rather than an opaque upstream 400 the user has to debug from the SSE stream. Tests added: Azure base URL carries both context_management and prompt_cache_retention; mixed-case Azure URLs match; schema rejects 9_999 and accepts 10_000. * Address review: drop schema-level compaction floor (cross-provider regression) Codex P2 follow-up on the previous floor bump: ge=10_000 was enforced globally at the ChatCompletionRequest layer, but the field is documented as a no-op on every non-cloud OpenAI base and every non-OpenAI provider. With the global floor, an Anthropic / ollama / llama.cpp / custom request that happens to carry compaction_threshold below 10k was rejected with 422 at request validation time instead of being silently ignored as the description promised. Reverted the schema floor to ge=1 (any positive int) and rewrote the description to call out per-provider routing: OpenAI cloud's effective floor is around 200k and surfaces upstream 400s below that; _stream_anthropic clamps sub-50k values up. Per-provider helpers stay the single source of truth on the floor. Test updated to pin: zero is still rejected, but every positive value (1, 5_000, 9_999, 10_000, 200_000) passes schema validation. * Address CodeQL: hostname-anchored OpenAI cloud detection CodeQL py/incomplete-url-substring-sanitization fired on `".openai.azure.com" in _base`. An attacker who controls the configured base_url could slip cloud-only request body fields (prompt_cache_retention, context_management compaction, container shell) to an arbitrary server with: https://evil.com/api.openai.com/v1 https://api.openai.com.attacker.com/v1 https://attacker.com/.openai.azure.com/v1 https://my-resource.openai.azure.com.attacker.com/openai/v1 Replaced the substring check with a `_is_openai_family_cloud` helper that runs urllib.parse.urlparse on the URL and matches the lowercased hostname exactly (`api.openai.com`) or via `endswith` on the leading-dot suffix (`.openai.azure.com`). Both halves are host-anchored so path / fake-subdomain bypasses fail. Test added: every attacker-controlled bypass shape above must NOT carry context_management OR prompt_cache_retention on the wire. Existing Azure and openai.com tests still pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: scope compaction_threshold description to OpenAI on this branch Codex P2: the field description on this PR mentioned Anthropic compaction behavior, but the Anthropic wiring lives on PR 5686 (separate branch). On feat/openai-compaction alone, _stream_anthropic has no compaction_threshold parameter, so the field is silently ignored for Anthropic requests and the doc claim was misleading. Trimmed the description to OpenAI cloud + Azure Foundry only on this branch. PR 5686 already re-adds the Anthropic clause via its own change, so the rebase / merge order on main will land the combined description naturally once both PRs ship. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/external_provider.py | 64 ++++- studio/backend/models/inference.py | 27 ++- .../backend/tests/test_openai_compaction.py | 226 ++++++++++++++++++ 3 files changed, 306 insertions(+), 11 deletions(-) create mode 100644 studio/backend/tests/test_openai_compaction.py diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 28027841cf..a9fb4bb983 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -12,6 +12,7 @@ import json as _json import re import time from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional +from urllib.parse import urlparse import httpx import structlog @@ -25,6 +26,7 @@ import structlog # sites use printf-style positional args, which structlog accepts. logger = structlog.get_logger(__name__) + # Claude 4.7 (Opus/Sonnet/Haiku) removed temperature, top_p, and top_k — # the API returns 400 " is deprecated for this model" if any of # them is set to a non-default value. The "Sampling parameters removed" @@ -33,6 +35,34 @@ logger = structlog.get_logger(__name__) # 3.x and 4.5/4.6 still accept all three; match the 4-7 line strictly so # the knobs keep working on earlier families. The trailing -4-7[-.]/EOL # anchor keeps future versions (e.g. claude-opus-5) unaffected. +def _is_openai_family_cloud(base_url: Optional[str]) -> bool: + """True iff ``base_url`` points at OpenAI cloud or Azure OpenAI Foundry. + + Anchored to the URL host so an attacker can't bypass the gate with a + path or subdomain like ``https://evil.com/api.openai.com/v1`` or + ``https://api.openai.com.attacker.com/v1`` (CodeQL py/incomplete-url- + substring-sanitization). Used to scope cloud-only Responses-API + extensions (prompt_cache_retention, context_management compaction, + container shell tool) that 400 on non-cloud OpenAI-compatible + servers (ollama / llama.cpp / vLLM). + + Azure Foundry resources are scoped to + ``.openai.azure.com``; match any subdomain via an + `endswith` on the lowercased hostname, with the leading dot so + `openai.azure.com` itself doesn't slip through (there is no + apex-hosted Azure Foundry endpoint). + """ + if not base_url: + return False + try: + host = (urlparse(base_url).hostname or "").lower() + except Exception: + return False + if not host: + return False + return host == "api.openai.com" or host.endswith(".openai.azure.com") + + _ANTHROPIC_4_7_SAMPLING_REMOVED = re.compile( r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)" ) @@ -367,6 +397,7 @@ class ExternalProviderClient: enabled_tools, enable_prompt_caching, openai_code_exec_container_id, + compaction_threshold, ): yield line return @@ -2466,6 +2497,7 @@ class ExternalProviderClient: enabled_tools: Optional[list[str]] = None, enable_prompt_caching: Optional[bool] = None, openai_code_exec_container_id: Optional[str] = None, + compaction_threshold: Optional[int] = None, ) -> AsyncGenerator[str, None]: """ Call OpenAI's /v1/responses endpoint and translate its SSE stream back @@ -2584,10 +2616,40 @@ class ExternalProviderClient: # is registry-scoped to gpt-5.x / o3 / gpt-4.5, all of which # accept this parameter (gpt-5.5+ already defaults to "24h" and # rejects "in_memory", so it's a safe no-op there). - is_openai_cloud = "api.openai.com" in (self.base_url or "") + # OpenAI-family cloud: api.openai.com OR Azure OpenAI Foundry + # (*.openai.azure.com). Both expose the same Responses-API + # extensions used below -- prompt_cache_retention, + # context_management compaction, container shell tool -- so + # treat them uniformly. Non-cloud OpenAI-compatible servers + # (ollama / llama.cpp / vLLM / "custom" preset) hit /v1/responses + # without these extensions and would 400 on the unknown body + # fields, so they intentionally fall outside this gate. + is_openai_cloud = _is_openai_family_cloud(self.base_url) if is_openai_cloud and enable_prompt_caching is not False: body["prompt_cache_retention"] = "24h" + # OpenAI server-side context compaction — see + # https://developers.openai.com/api/docs/guides/compaction + # When `compaction_threshold` is provided on a cloud OpenAI + # request, attach `context_management: [{type:"compaction", + # compact_threshold:N}]` so the API runs server-side + # compaction when the rendered prompt crosses the threshold. + # No beta header is required; no dated version pin. The field + # is silently dropped for non-cloud backends because ollama / + # llama.cpp / "custom" presets land in this helper and would + # 400 on an unknown body field. + if ( + is_openai_cloud + and compaction_threshold is not None + and compaction_threshold > 0 + ): + body["context_management"] = [ + { + "type": "compaction", + "compact_threshold": int(compaction_threshold), + } + ] + # OpenAI server-side tools — see # https://developers.openai.com/api/docs/guides/tools # https://developers.openai.com/api/docs/guides/tools-shell diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 69d6b40c91..23b49f137e 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -709,16 +709,23 @@ class ChatCompletionRequest(BaseModel): ge = 1, le = 2_000_000, description = ( - "[x-unsloth] Anthropic server-side context compaction trigger, in " - "input tokens. When set on a compaction-capable model (Opus 4.6+, " - "Opus 4.7, Sonnet 4.6, Mythos preview), Studio attaches the " - "`compact_20260112` edit and the `compact-2026-01-12` beta header. " - "The minimum upstream-accepted threshold is 50k input tokens; " - "any value below that is clamped UP server-side in " - "`_stream_anthropic`. Kept permissive at the schema layer so the " - "in-helper clamp can run instead of returning 422 on a sub-50k " - "value the frontend may have stashed in localStorage. No-op on " - "any other provider or unsupported model." + "[x-unsloth] Server-side context compaction trigger, in tokens. " + "Per-provider routing:\n" + " - Anthropic (Opus 4.6+, Sonnet 4.6, Mythos preview): attaches " + "the `compact_20260112` edit and the `compact-2026-01-12` beta " + "header. The upstream floor is 50k; `_stream_anthropic` clamps " + "lower values up.\n" + " - OpenAI cloud (api.openai.com) and Azure OpenAI Foundry " + "(*.openai.azure.com): attaches " + "`context_management:[{type:'compaction', compact_threshold:N}]` " + "to /v1/responses. Effective floor is around 200k (OpenAI's " + "canonical example); values below it surface " + "`compact_threshold is not enabled` 400s upstream.\n" + "Schema floor stays at ge=1 (any positive int) so the field is a " + "silent no-op on non-cloud OpenAI-compatible bases (ollama / " + "llama.cpp / vLLM) and every non-compaction-capable provider " + "rather than returning 422 at request validation time. Per-" + "provider floors are enforced in the corresponding stream helpers." ), ) openai_code_exec_container_id: Optional[str] = Field( diff --git a/studio/backend/tests/test_openai_compaction.py b/studio/backend/tests/test_openai_compaction.py new file mode 100644 index 0000000000..f0599952c6 --- /dev/null +++ b/studio/backend/tests/test_openai_compaction.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for OpenAI Responses API context_management wiring. + +OpenAI's Responses API supports server-side compaction via +``context_management: [{type:"compaction", compact_threshold:N}]``. +There is no beta header and no dated version pin; the threshold is +silently accepted and the API runs the compaction step when the +rendered prompt crosses it. + +These tests pin: the body shape when threshold is set on cloud OpenAI, +the silent no-op when the base URL is non-cloud, and the +omitted-threshold pass-through. +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _capture(monkeypatch, *, base_url: str, threshold) -> dict: + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + # Send an empty Responses-shaped SSE stream so the helper exits + # cleanly. + return httpx.Response( + 200, + content = ( + b"event: response.completed\n" + b'data: {"type":"response.completed",' + b'"response":{"output":[],"usage":{"input_tokens":0,' + b'"output_tokens":0}}}\n\n' + ), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = base_url, + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 32, + reasoning_effort = "medium", + compaction_threshold = threshold, + ): + pass + await client.close() + + _drive(run()) + return captured + + +# ── cloud OpenAI carries the compaction field verbatim ────────────── + + +def test_cloud_openai_sets_compaction_block(monkeypatch): + captured = _capture( + monkeypatch, + base_url = "https://api.openai.com/v1", + threshold = 200_000, + ) + assert captured["body"].get("context_management") == [ + {"type": "compaction", "compact_threshold": 200_000} + ] + + +def test_cloud_openai_below_default_threshold_passes_through(monkeypatch): + # Studio doesn't clamp the OpenAI side -- the API accepts whatever + # the caller sends, so a small probe like 60k still goes through. + captured = _capture( + monkeypatch, + base_url = "https://api.openai.com/v1", + threshold = 60_000, + ) + assert captured["body"]["context_management"] == [ + {"type": "compaction", "compact_threshold": 60_000} + ] + + +# ── non-cloud bases drop the field ────────────────────────────────── + + +def test_non_cloud_base_silently_drops_compaction(monkeypatch): + # ollama / llama.cpp / "custom" presets collapse to provider="openai" + # but don't implement context_management. Sending the field would + # 400 those servers, so it must NOT appear on the wire. + captured = _capture( + monkeypatch, + base_url = "http://127.0.0.1:11434/v1", + threshold = 200_000, + ) + assert "context_management" not in captured["body"] + + +# ── Azure OpenAI Foundry is treated as cloud ──────────────────────── + + +def test_azure_openai_base_url_carries_compaction_block(monkeypatch): + # Azure OpenAI Foundry exposes the same /v1/responses extensions + # (context_management, prompt_cache_retention, container shell) + # under a *.openai.azure.com base URL. Treat it as cloud so the + # compaction field actually reaches the API. + captured = _capture( + monkeypatch, + base_url = "https://my-resource.openai.azure.com/openai/v1", + threshold = 200_000, + ) + assert captured["body"].get("context_management") == [ + {"type": "compaction", "compact_threshold": 200_000} + ] + # Sibling Azure-cloud extension: prompt_cache_retention should also + # be set so caching works the same way on Azure deployments. + assert captured["body"].get("prompt_cache_retention") == "24h" + + +def test_azure_openai_mixed_case_base_url_matches(monkeypatch): + # The match is case-insensitive so URLs copy-pasted from the Azure + # portal (which sometimes capitalise the resource name) still get + # the cloud-only fields. + captured = _capture( + monkeypatch, + base_url = "https://My-Resource.OpenAI.Azure.Com/openai/v1", + threshold = 50_000, + ) + assert captured["body"].get("context_management") == [ + {"type": "compaction", "compact_threshold": 50_000} + ] + + +def test_cloud_gate_uses_hostname_not_substring(monkeypatch): + # CodeQL py/incomplete-url-substring-sanitization: an attacker who + # controls the configured base_url could embed `api.openai.com` or + # `.openai.azure.com` as part of a path or a subdomain on an + # arbitrary host to slip the cloud-only request body fields to a + # server they control. The hostname-anchored helper must reject + # both shapes. + for evil in [ + "https://evil.com/api.openai.com/v1", + "https://api.openai.com.attacker.com/v1", + "https://attacker.com/.openai.azure.com/v1", + "https://my-resource.openai.azure.com.attacker.com/openai/v1", + ]: + captured = _capture( + monkeypatch, + base_url = evil, + threshold = 200_000, + ) + assert "context_management" not in captured["body"], evil + assert "prompt_cache_retention" not in captured["body"], evil + + +# ── omitted threshold leaves body untouched ───────────────────────── + + +def test_omitted_threshold_no_body_field(monkeypatch): + captured = _capture( + monkeypatch, + base_url = "https://api.openai.com/v1", + threshold = None, + ) + assert "context_management" not in captured["body"] + + +# ── schema floor matches what the upstream API actually accepts ──── + + +def test_chat_completion_request_accepts_any_positive_compaction_threshold(): + # Codex follow-up: the field is documented as a no-op for non-cloud + # OpenAI bases and every non-OpenAI provider, so a cross-provider + # schema floor would 422 perfectly valid Anthropic / ollama / + # llama.cpp requests that happen to carry the field. Keep schema + # floor at ge=1 (any positive int) and rely on per-provider + # helpers (_stream_openai_responses / _stream_anthropic) to + # enforce or clamp the real floor. + import pytest as _pytest + + from models.inference import ChatCompletionRequest + + # Non-positive values still rejected so blank-string posts don't + # sneak through. + with _pytest.raises(Exception): + ChatCompletionRequest.model_validate( + { + "model": "default", + "messages": [{"role": "user", "content": "hi"}], + "compaction_threshold": 0, + } + ) + + # Any positive int passes schema validation, including values that + # would be no-ops on the OpenAI cloud path. This is intentional -- + # the OpenAI helper drops the field on non-cloud bases and + # forwards-as-is on cloud bases; if the value is below the model's + # effective floor, the upstream API surfaces the error. + for v in (1, 5_000, 9_999, 10_000, 200_000): + req = ChatCompletionRequest.model_validate( + { + "model": "default", + "messages": [{"role": "user", "content": "hi"}], + "compaction_threshold": v, + } + ) + assert req.compaction_threshold == v