From 39706696fc515727aba056a1850082270d736443 Mon Sep 17 00:00:00 2001 From: Mitali Agrawal Date: Tue, 25 Aug 2026 16:42:10 -0700 Subject: [PATCH] refactor(server): dedupe maintenance gate across API entrypoints Extract shared maintenance_state checks into maintenance.py so OpenAI, Anthropic, Responses, and /generate all return consistent 503 messages. Co-authored-by: Cursor --- python/freetoken/server/anthropic_api.py | 5 ++-- python/freetoken/server/api_server.py | 6 ++--- python/freetoken/server/maintenance.py | 29 ++++++++++++++++++++++++ python/freetoken/server/openai_api.py | 21 +++-------------- python/freetoken/server/responses_api.py | 5 ++-- tests/server/test_rebuild_maintenance.py | 12 +++++----- 6 files changed, 45 insertions(+), 33 deletions(-) create mode 100644 python/freetoken/server/maintenance.py diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index c941d98b..d7a4d547 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -52,6 +52,7 @@ submit_generation, with_keepalive, ) +from .maintenance import maintenance_state_of, maintenance_unavailable_detail from .request_logger import log_request # Emit a protocol-native `ping` event after this many seconds of stream silence, @@ -83,9 +84,7 @@ def register_anthropic_routes( async def v1_messages(req: AnthropicMessagesRequest, request: Request): log_request("/v1/messages", req, request) state = get_state() - mstate = getattr(state, "maintenance_state", "serving") - if mstate != "serving": - detail = "model is still loading" if mstate == "loading" else "cache rebuild in progress" + if (detail := maintenance_unavailable_detail(maintenance_state_of(state))) is not None: return _anthropic_error_response(503, "overloaded_error", detail) return await handle_anthropic_messages(req, request, state, get_model_sampling()) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc85..81654066 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -39,6 +39,7 @@ from .anthropic_api import register_anthropic_routes from .accounting import AdmissionClosedError, register_accounting_routes from .control_api import register_control_routes +from .maintenance import maintenance_gate from .openai_api import register_openai_routes from . import request_ring from .access_log_filter import install_polling_access_log_filter @@ -822,9 +823,8 @@ async def generate(req: GenerateRequest, request: Request): logger.debug("Received generate request %s", req) log_request("/generate", req, request) state = get_global_state() - if state.maintenance_state != "serving": - detail = "model is still loading" if state.maintenance_state == "loading" else "cache rebuild in progress" - return JSONResponse({"error": f"server unavailable: {detail}"}, status_code=503) + if (gate := maintenance_gate(state)) is not None: + return gate if req.max_tokens < 1: return JSONResponse({"error": f"max_tokens must be at least 1, got {req.max_tokens}"}, status_code=400) uid = state.new_user() diff --git a/python/freetoken/server/maintenance.py b/python/freetoken/server/maintenance.py new file mode 100644 index 00000000..41d7179a --- /dev/null +++ b/python/freetoken/server/maintenance.py @@ -0,0 +1,29 @@ +"""Shared maintenance-gate helpers for API entrypoints.""" + +from __future__ import annotations + +from typing import Any + +from fastapi.responses import JSONResponse + + +def maintenance_state_of(state: Any) -> str: + return getattr(state, "maintenance_state", "serving") + + +def maintenance_unavailable_detail(mstate: str) -> str | None: + """Client-facing unavailable message, or None when serving.""" + if mstate == "serving": + return None + if mstate == "loading": + return "model is still loading" + if mstate == "failed": + return "server unavailable: maintenance failed (restart required)" + return "server unavailable: cache rebuild in progress" + + +def maintenance_gate(state: Any) -> JSONResponse | None: + """503 while the engine is not serving. None when serving.""" + if (msg := maintenance_unavailable_detail(maintenance_state_of(state))) is None: + return None + return JSONResponse({"error": msg}, status_code=503) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd26..71cf446e 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -21,6 +21,7 @@ ToolChoiceObject, ) from .function_call_parser import ToolCallItem +from .maintenance import maintenance_gate from .request_logger import log_request from .generation import ( ContentDelta, @@ -87,22 +88,6 @@ def _all_tool_dicts(tools) -> list[dict[str, Any]]: return [t.model_dump(exclude_none=True) for t in (tools or [])] -def _maintenance_gate(state: Any) -> JSONResponse | None: - """503 while the engine is not serving. Distinguishes the startup "loading" phase from a - runtime cache "rebuild"/"failed" so clients (and the desktop) get an actionable message. - None when serving.""" - mstate = getattr(state, "maintenance_state", "serving") - if mstate == "serving": - return None - if mstate == "loading": - msg = "model is still loading" - elif mstate == "failed": - msg = "server unavailable: maintenance failed (restart required)" - else: - msg = "server unavailable: cache rebuild in progress" - return JSONResponse({"error": msg}, status_code=503) - - def register_openai_routes( app: FastAPI, get_state: Callable[[], Any], @@ -116,7 +101,7 @@ async def v1_root(): async def v1_chat_completions(req: ChatCompletionRequest, request: Request): log_request("/v1/chat/completions", req, request) state = get_state() - if (gate := _maintenance_gate(state)) is not None: + if (gate := maintenance_gate(state)) is not None: return gate return await handle_chat_completion(req, request, state, get_model_sampling()) @@ -124,7 +109,7 @@ async def v1_chat_completions(req: ChatCompletionRequest, request: Request): async def v1_completions(req: CompletionRequest, request: Request): log_request("/v1/completions", req, request) state = get_state() - if (gate := _maintenance_gate(state)) is not None: + if (gate := maintenance_gate(state)) is not None: return gate return await handle_completion(req, request, state, get_model_sampling()) diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index 60a1442f..a0f47a75 100644 --- a/python/freetoken/server/responses_api.py +++ b/python/freetoken/server/responses_api.py @@ -77,6 +77,7 @@ submit_generation, with_keepalive, ) +from .maintenance import maintenance_state_of, maintenance_unavailable_detail from .request_logger import log_request # Seconds of event silence before a keep-alive frame is emitted on the stream. @@ -117,9 +118,7 @@ def register_responses_routes( async def v1_responses(req: ResponsesRequest, request: Request): log_request("/v1/responses", req, request) state = get_state() - mstate = getattr(state, "maintenance_state", "serving") - if mstate != "serving": - detail = "model is still loading" if mstate == "loading" else "cache rebuild in progress" + if (detail := maintenance_unavailable_detail(maintenance_state_of(state))) is not None: return _error_response(503, detail) if req.background: return _error_response(400, "background mode is not supported") diff --git a/tests/server/test_rebuild_maintenance.py b/tests/server/test_rebuild_maintenance.py index a1d10a7e..1d39695b 100644 --- a/tests/server/test_rebuild_maintenance.py +++ b/tests/server/test_rebuild_maintenance.py @@ -226,19 +226,19 @@ def on_failure(message): def test_openai_gate_message_is_loading_aware(): - from freetoken.server.openai_api import _maintenance_gate + from freetoken.server.maintenance import maintenance_gate - assert _maintenance_gate(SimpleNamespace(maintenance_state="serving")) is None - loading = _maintenance_gate(SimpleNamespace(maintenance_state="loading")) + assert maintenance_gate(SimpleNamespace(maintenance_state="serving")) is None + loading = maintenance_gate(SimpleNamespace(maintenance_state="loading")) assert loading is not None and loading.status_code == 503 assert b"loading" in loading.body.lower() - rebuild = _maintenance_gate(SimpleNamespace(maintenance_state="rebuilding")) + rebuild = maintenance_gate(SimpleNamespace(maintenance_state="rebuilding")) assert rebuild is not None and rebuild.status_code == 503 assert b"rebuild" in rebuild.body.lower() - failed = _maintenance_gate(SimpleNamespace(maintenance_state="failed")) + failed = maintenance_gate(SimpleNamespace(maintenance_state="failed")) assert failed is not None and failed.status_code == 503 # A state object without the attribute defaults to serving (defensive, never blocks). - assert _maintenance_gate(SimpleNamespace()) is None + assert maintenance_gate(SimpleNamespace()) is None def test_cache_rebuild_guarded_during_loading():