From 494aa623a50e238d3e0244f2b859a2df4e96c322 Mon Sep 17 00:00:00 2001 From: mac Date: Wed, 26 Aug 2026 23:15:01 +0300 Subject: [PATCH] server: validate sampling params and rebuild timeout at the API layer resolve_sampling forwarded arbitrary client floats into SamplingParams: negative, NaN and inf temperature, top_p outside (0, 1], top_k 0 -- all reached the engine and died (or misbehaved) deep in the sampler instead of failing the request. Validate the RESOLVED values (checkpoint defaults included) and raise ValueError in the style of the existing max_tokens check; every adapter already maps that to a clean 400. CacheRebuildRequest.timeout gets the matching pydantic bound (0 < t <= 3600) so a NaN/inf/negative timeout cannot reach asyncio.wait_for. Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defects 3-4); the rebuild TOCTOU race from that audit is already structurally fixed upstream, so only the timeout bound is ported. --- python/freetoken/server/api_server.py | 4 +- python/freetoken/server/generation.py | 16 ++++-- tests/server/test_openai_api.py | 63 ++++++++++++++++++++++++ tests/server/test_rebuild_maintenance.py | 15 ++++++ 4 files changed, 93 insertions(+), 5 deletions(-) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc854..94b3e10ef 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -33,7 +33,7 @@ init_logger, load_generation_sampling, ) -from pydantic import BaseModel +from pydantic import BaseModel, Field from .args import ServerArgs from .anthropic_api import register_anthropic_routes @@ -494,7 +494,7 @@ class CacheRebuildRequest(BaseModel): # is deferred (needs the drain-gate machinery); constraining the Literal makes an # unsupported value fail fast with a 422 at the API layer instead of a generic 503. mode: Literal["if_idle"] = "if_idle" - timeout: float = 300.0 + timeout: float = Field(default=300.0, gt=0, le=3600) async def dispatch_rebuild( diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908a..a60119ecf 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -15,6 +15,7 @@ import asyncio import json +import math import time from collections.abc import AsyncIterator from dataclasses import dataclass, field @@ -177,12 +178,21 @@ def pick(value, key, framework): # non-positive value is a client error. if max_tokens is not None and max_tokens < 1: raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") + resolved_temperature = pick(temperature, "temperature", 0.0) + if not math.isfinite(resolved_temperature) or resolved_temperature < 0: + raise ValueError(f"temperature must be a finite number >= 0, got {resolved_temperature}") + resolved_top_p = pick(top_p, "top_p", 1.0) + if not math.isfinite(resolved_top_p) or not 0 < resolved_top_p <= 1: + raise ValueError(f"top_p must be in (0, 1], got {resolved_top_p}") + resolved_top_k = pick(top_k, "top_k", -1) + if resolved_top_k != -1 and resolved_top_k < 1: + raise ValueError(f"top_k must be -1 (disabled) or >= 1, got {resolved_top_k}") return SamplingParams( ignore_eos=ignore_eos, max_tokens=DEFAULT_MAX_OUTPUT_TOKENS if max_tokens is None else max_tokens, - temperature=pick(temperature, "temperature", 0.0), - top_k=pick(top_k, "top_k", -1), - top_p=pick(top_p, "top_p", 1.0), + temperature=resolved_temperature, + top_k=resolved_top_k, + top_p=resolved_top_p, stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) ) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index facd469b3..f6808a7d3 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -676,3 +676,66 @@ def test_minimax_http_non_stream_forces_implicit_reasoning_without_request_knob( message = response["choices"][0]["message"] assert message["reasoning_content"] == "private thought" assert message["content"] == "visible answer" + + +import pytest # noqa: E402 + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("temperature", -0.5), + ("top_p", 0.0), + ("top_k", 0), + ], +) +def test_chat_completion_rejects_invalid_sampling(field, value): + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + body = { + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + field: value, + } + response = TestClient(app).post("/v1/chat/completions", json=body) + + assert response.status_code == 400 + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert field in error["message"] + assert state.sent is None + + +@pytest.mark.parametrize( + "sampling", + [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1}, + {"temperature": 0.0, "top_p": 1.0, "top_k": 1}, + ], +) +def test_chat_completion_accepts_sampling_boundaries(sampling): + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + response = TestClient(app).post( + "/v1/chat/completions", + json={ + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + **sampling, + }, + ) + + assert response.status_code == 200 + assert state.sent is not None diff --git a/tests/server/test_rebuild_maintenance.py b/tests/server/test_rebuild_maintenance.py index a1d10a7e5..0df3d6bfe 100644 --- a/tests/server/test_rebuild_maintenance.py +++ b/tests/server/test_rebuild_maintenance.py @@ -327,3 +327,18 @@ def test_cache_rebuild_request_rejects_unknown_mode(): assert CacheRebuildRequest(mode="if_idle").mode == "if_idle" with pytest.raises(ValidationError): CacheRebuildRequest(mode="drain") + + +from pydantic import ValidationError # noqa: E402 + +from freetoken.server.api_server import CacheRebuildRequest # noqa: E402 + + +def test_rebuild_timeout_bounds(): + import pytest + + for bad in (0, -5, 3601, float("nan")): + with pytest.raises(ValidationError): + CacheRebuildRequest(timeout=bad) + assert CacheRebuildRequest(timeout=3600).timeout == 3600 + assert CacheRebuildRequest().timeout == 300.0