diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 5592a2b0d..7267b1a9c 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -710,6 +710,30 @@ def _reasoning_geometry(state: Any) -> dict | None: return {"gears": list(gears), "default": default, "kwargs": kwargs} +def kv_pool_geometry(state: Any) -> tuple[int, int]: + """The allocated KV pool as ``(num_pages, page_size)``, most-recent-truth first: the last + rebuild's result, else the running UserReply snapshot, else the load-time ("meta", …) ack. + ``num_pages`` is 0 until one of those lands (still loading / older engine build). + + Split out of cache_geometry so /v1/models can clamp the advertised context to the pool + without recomputing the whole panel payload -- both callers must agree on the resolution + order or they would report different capacities for the same server. + """ + stats = getattr(state, "stats", None) + last = getattr(state, "last_rebuild", None) or {} + pools = getattr(state, "cache_pools", None) or {} + num_pages = int( + last.get("num_pages") + or getattr(stats, "kv_total_pages", 0) + or pools.get("num_pages", 0) + or 0 + ) + page_size = int( + pools.get("page_size", 0) or getattr(getattr(state, "config", None), "page_size", 1) or 1 + ) + return num_pages, page_size + + def cache_geometry(state: Any) -> dict: """Current cache geometry for the desktop cache panel. Each pool size resolves most-recent-truth first: the last rebuild's result, else the running UserReply snapshot @@ -725,11 +749,10 @@ def cache_geometry(state: Any) -> dict: config = state.config last = getattr(state, "last_rebuild", None) or {} pools = getattr(state, "cache_pools", None) or {} - num_pages = int(last.get("num_pages") or tr.kv_total_pages or pools.get("num_pages", 0) or 0) + num_pages, page_size = kv_pool_geometry(state) num_mamba_slots = int( last.get("mamba_slots") or tr.mamba_total_slots or pools.get("num_mamba_slots", 0) or 0 ) - page_size = int(pools.get("page_size", 0) or getattr(config, "page_size", 1) or 1) moe_cache_size = last.get("moe_cache_size") if moe_cache_size is None: moe_cache_size = int(pools.get("moe_cache_size", 0) or 0) or configured_moe_cache_size( diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index dc2f73a97..087b2a4e4 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -701,10 +701,31 @@ def _served_model_name(state: Any) -> str: def _model_context_length(state: Any) -> int | None: - """The model ceiling, not `min(ceiling, KV budget)`: a rebuild moves the latter, and agents - read this once at startup.""" + """The usable context: ``min(model ceiling, allocated KV tokens)``. + + This is the clamp the engine already applies to its own ``max_seq_len`` (Engine.__init__, + and again in _refresh_seq_state after a rebuild), and therefore the limit the scheduler + admits against. The frontend process holds an unclamped ServerArgs copy, so publishing + ``config.max_seq_len`` here advertised the model's positional ceiling instead -- e.g. 262144 + against a 178176-token pool, so the engine and this route disagreed about one quantity. + + `ft launch` reads this to size each agent's context window (opencode's ``limit.context``, + codex's ``context_window``, ``CLAUDE_CODE_MAX_CONTEXT_TOKENS``), so overstating it stops + those agents compacting before the KV pool runs out. + """ try: # never 500 a metadata route: max_seq_len walks into the HF config on some builds value = int(state.config.max_seq_len) except Exception: # noqa: BLE001 return None - return value if value > 0 else None + if value <= 0: + return None + try: + # Local import: api_server imports this module, so a module-level one would cycle. + from .api_server import kv_pool_geometry + + num_pages, page_size = kv_pool_geometry(state) + kv_tokens = num_pages * page_size + except Exception: # noqa: BLE001 + kv_tokens = 0 + # 0 before the ("meta", …) ack lands; the ceiling is the best available answer until then. + return min(value, kv_tokens) if kv_tokens > 0 else value diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index e33018e19..d9d565ff6 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -442,19 +442,67 @@ def test_models_route_returns_served_model_name(): assert card["max_model_len"] is None and card["context_length"] is None +def _models_card(state): + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + return TestClient(app).get("/v1/models").json()["data"][0] + + def test_models_route_publishes_the_model_context_length(): """`ft launch` reads this to size each agent's context window.""" state = FakeState([]) state.config.max_seq_len = 262144 - app = FastAPI() - register_openai_routes(app, lambda: state, lambda: {}) - card = TestClient(app).get("/v1/models").json()["data"][0] + card = _models_card(state) + # No pool geometry resolved yet (still loading, or pre-("meta", …) ack): the model + # ceiling is the only answer available. assert card["max_model_len"] == 262144 assert card["context_length"] == 262144 +def test_models_route_clamps_the_context_length_to_the_kv_pool(): + """The engine clamps its own max_seq_len to the allocated pool and the scheduler admits + against that, so publishing the unclamped ceiling here made `ft launch` size each agent's + compaction window past what the server can actually hold (#448).""" + state = FakeState([]) + state.config.max_seq_len = 262144 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + + card = _models_card(state) + + assert card["max_model_len"] == 178176 + assert card["context_length"] == 178176 + + +def test_models_route_keeps_the_ceiling_when_the_pool_exceeds_it(): + """A pool larger than the model's positional ceiling does not extend the context.""" + state = FakeState([]) + state.config.max_seq_len = 32768 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + + card = _models_card(state) + + assert card["max_model_len"] == 32768 + assert card["context_length"] == 32768 + + +def test_models_route_prefers_the_last_rebuild_over_the_load_time_pool(): + """A rebuild moves the pool; /v1/models must follow it rather than freeze the load-time + allocation (same most-recent-truth order /v1/cache/status reports).""" + state = FakeState([]) + state.config.max_seq_len = 262144 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + state.last_rebuild = {"num_pages": 40000} + + card = _models_card(state) + + assert card["max_model_len"] == 40000 + + async def _collect(generator): return [chunk async for chunk in generator]