ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: wire OpenAI Responses server-side context compaction (#5687)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-05-22 06:20:45 -07:00 • committed by GitHub
commit e86f3c5dc7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 306 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

64
studio/backend/core/inference/external_provider.py
View file

@ -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 "<param> 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
``<resource-name>.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

27
studio/backend/models/inference.py
View file

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

226
studio/backend/tests/test_openai_compaction.py Normal file
View file

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