mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-26 05:44:46 -06:00
chore(coord): remove spawn-quota subsystem (#403)
* chore(coord): remove spawn-quota subsystem The quota gate was operator-level safety per its own comments, not a security boundary, and never fired in a week of heavy use. Runaway coordinator spawns are already bounded by max_active slot exhaustion, which surfaces to the coord LLM as a tool error — same operational shape, one fewer moving part. Precedes the Stage 1 SessionManager unification so the coord tool doesn't inherit quota bookkeeping. Upgraded deployments with the three removed settings persisted will log three "Skipping invalid setting" warnings on startup and otherwise degrade cleanly; a follow-up migration to delete the rows would silence that noise. * chore(migrations): drop stale coord spawn-quota settings rows (047) Clears persisted rows for the three ConfigStore keys removed in the previous commit so upgraded deployments don't log "Skipping invalid setting" warnings on every startup. Downgrade is a no-op — the rows were operator-set values, and a rollback to pre-1.5.0 code falls back to the registry defaults for any key not present.
This commit is contained in:
Patrick Buckley
@@ -2,7 +2,7 @@
|
||||
"openapi": "3.1.0",
|
||||
"info": {
|
||||
"title": "turnstone Console API",
|
||||
"version": "1.5.0a2",
|
||||
"version": "1.5.0a4",
|
||||
"description": "Cluster-wide visibility and control across all turnstone nodes."
|
||||
},
|
||||
"paths": {
|
||||
@@ -5732,158 +5732,6 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"/v1/api/coordinator/{ws_id}/quota": {
|
||||
"get": {
|
||||
"summary": "Read the coordinator session's live spawn-quota state",
|
||||
"operationId": "v1_api_coordinator_{ws_id}_quota_get",
|
||||
"tags": [
|
||||
"Coordinator"
|
||||
],
|
||||
"description": "Returns the current ``spawn_budget`` (active-children cap) and ``spawn_rate`` bucket (``tokens_per_minute``, ``burst``, ``tokens_available``). Values reflect the in-memory override when an admin has mutated the session via POST; otherwise they reflect the global defaults baked in at session construction.",
|
||||
"parameters": [
|
||||
{
|
||||
"name": "ws_id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Success",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/CoordinatorQuotaResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"400": {
|
||||
"description": "Error 400",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"403": {
|
||||
"description": "Error 403",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"404": {
|
||||
"description": "Error 404",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"503": {
|
||||
"description": "Error 503",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"post": {
|
||||
"summary": "Mutate the coordinator session's spawn quota (partial update)",
|
||||
"operationId": "v1_api_coordinator_{ws_id}_quota_post",
|
||||
"tags": [
|
||||
"Coordinator"
|
||||
],
|
||||
"description": "Updates any subset of ``spawn_budget``, ``spawn_rate.tokens_per_minute``, and ``spawn_rate.burst``. Nested ``spawn_rate`` and flat ``tokens_per_minute`` / ``burst`` aliases are both accepted. Missing fields keep their current values. Overrides are in-memory only \u2014 a session reopen re-seeds from the global settings, matching the /trust and /restrict contract. Writes ``coordinator.quota.updated`` with the before/after snapshot.",
|
||||
"parameters": [
|
||||
{
|
||||
"name": "ws_id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/CoordinatorQuotaRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Success",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/CoordinatorQuotaResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"400": {
|
||||
"description": "Error 400",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"403": {
|
||||
"description": "Error 403",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"404": {
|
||||
"description": "Error 404",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"503": {
|
||||
"description": "Error 503",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/v1/api/cluster/ws/{ws_id}/detail": {
|
||||
"get": {
|
||||
"summary": "Cluster-wide live workstream detail (storage + live block + tail)",
|
||||
@@ -7476,182 +7324,6 @@
|
||||
"title": "CoordinatorOpenResponse",
|
||||
"type": "object"
|
||||
},
|
||||
"CoordinatorQuotaRequest": {
|
||||
"description": "Body for POST /v1/api/coordinator/{ws_id}/quota.\n\nPartial-update semantics \u2014 any subset of the three knobs may be\nsupplied; missing fields keep their current values. Accepts either\nthe nested ``spawn_rate`` object OR the flat ``tokens_per_minute``\n/ ``burst`` aliases \u2014 supplying both for the same field yields a\n400 so the admin UI can't half-migrate its body shape unnoticed.",
|
||||
"properties": {
|
||||
"spawn_budget": {
|
||||
"anyOf": [
|
||||
{
|
||||
"maximum": 500,
|
||||
"minimum": 1,
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "New active-children cap (1..500).",
|
||||
"title": "Spawn Budget"
|
||||
},
|
||||
"spawn_rate": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/CoordinatorSpawnRateInput"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "Nested rate-bucket overrides. Use this OR the flat tokens_per_minute/burst aliases for a given field; mixing both shapes for the same field is rejected with 400."
|
||||
},
|
||||
"tokens_per_minute": {
|
||||
"anyOf": [
|
||||
{
|
||||
"maximum": 600.0,
|
||||
"minimum": 0.0,
|
||||
"type": "number"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "Flat alias for spawn_rate.tokens_per_minute.",
|
||||
"title": "Tokens Per Minute"
|
||||
},
|
||||
"burst": {
|
||||
"anyOf": [
|
||||
{
|
||||
"maximum": 500,
|
||||
"minimum": 1,
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "Flat alias for spawn_rate.burst.",
|
||||
"title": "Burst"
|
||||
}
|
||||
},
|
||||
"title": "CoordinatorQuotaRequest",
|
||||
"type": "object"
|
||||
},
|
||||
"CoordinatorSpawnRateInput": {
|
||||
"description": "Nested rate-bucket sub-object of the quota REQUEST body.\n\nIntentionally excludes the response-only ``tokens_available`` field\nso generated SDK input types don't imply clients can POST a\nlive-bucket reading \u2014 the server ignores it on input.",
|
||||
"properties": {
|
||||
"tokens_per_minute": {
|
||||
"anyOf": [
|
||||
{
|
||||
"maximum": 600.0,
|
||||
"minimum": 0.0,
|
||||
"type": "number"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "New refill rate in tokens per minute.",
|
||||
"title": "Tokens Per Minute"
|
||||
},
|
||||
"burst": {
|
||||
"anyOf": [
|
||||
{
|
||||
"maximum": 500,
|
||||
"minimum": 1,
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "New burst ceiling.",
|
||||
"title": "Burst"
|
||||
}
|
||||
},
|
||||
"title": "CoordinatorSpawnRateInput",
|
||||
"type": "object"
|
||||
},
|
||||
"CoordinatorQuotaResponse": {
|
||||
"description": "Response body for GET/POST /v1/api/coordinator/{ws_id}/quota.",
|
||||
"properties": {
|
||||
"status": {
|
||||
"default": "ok",
|
||||
"title": "Status",
|
||||
"type": "string"
|
||||
},
|
||||
"spawn_budget": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "Active-children cap. None on a non-coordinator session.",
|
||||
"title": "Spawn Budget"
|
||||
},
|
||||
"spawn_rate": {
|
||||
"$ref": "#/components/schemas/CoordinatorSpawnRateState",
|
||||
"description": "Rate-bucket state. Fields are None on non-coordinator sessions."
|
||||
}
|
||||
},
|
||||
"title": "CoordinatorQuotaResponse",
|
||||
"type": "object"
|
||||
},
|
||||
"CoordinatorSpawnRateState": {
|
||||
"description": "Nested rate-bucket sub-object of the quota RESPONSE body.\n\nExtends the input shape with the read-only ``tokens_available``\nsnapshot so the admin UI can render a \"rate status\" badge.",
|
||||
"properties": {
|
||||
"tokens_per_minute": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "number"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "Current refill rate (tokens/minute). None on a non-coordinator session.",
|
||||
"title": "Tokens Per Minute"
|
||||
},
|
||||
"burst": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "Current burst ceiling. None on a non-coordinator session.",
|
||||
"title": "Burst"
|
||||
},
|
||||
"tokens_available": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "number"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"description": "Current post-refill token count \u2014 advisory only; the bucket refills continuously so this snapshot is stale the moment it's read. Useful for the admin UI's 'rate status' badge. Response-only: the server ignores this field on input.",
|
||||
"title": "Tokens Available"
|
||||
}
|
||||
},
|
||||
"title": "CoordinatorSpawnRateState",
|
||||
"type": "object"
|
||||
},
|
||||
"CoordinatorRestrictRequest": {
|
||||
"description": "Body for POST /v1/api/coordinator/{ws_id}/restrict.",
|
||||
"properties": {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"openapi": "3.1.0",
|
||||
"info": {
|
||||
"title": "turnstone Server API",
|
||||
"version": "1.5.0a2",
|
||||
"version": "1.5.0a4",
|
||||
"description": "Single-node workstream management, chat interaction, and real-time streaming."
|
||||
},
|
||||
"paths": {
|
||||
|
||||
@@ -481,66 +481,6 @@ def test_list_children_skill_filter_avoids_n_plus_one(populated_storage, monkeyp
|
||||
assert call_count["n"] == 0
|
||||
|
||||
|
||||
def test_count_active_children_counts_non_terminal_states(populated_storage):
|
||||
"""Budget count must use an aggregate SQL query so a tail of
|
||||
recently-closed children can't push live rows past a LIMIT and
|
||||
silently undercount (Copilot #3 on PR #387).
|
||||
|
||||
populated_storage has:
|
||||
- coord-1 (coordinator, excluded)
|
||||
- child-a (interactive, idle) → counted
|
||||
- child-b (interactive, running) → counted
|
||||
- child-coord (coordinator child) → excluded (kind filter doesn't apply
|
||||
to count_workstreams_by_state, but
|
||||
it still matches parent_ws_id+user_id)
|
||||
- unrelated (no parent) → excluded (parent filter)
|
||||
- cross-tenant-child (user-2) → excluded (user_id filter)
|
||||
|
||||
child-coord DOES count against count_workstreams_by_state because
|
||||
the aggregate doesn't filter by kind — the budget is per-coord
|
||||
across any descendant type. That's fine semantically: a
|
||||
coordinator that spawns a nested coord still occupies a slot.
|
||||
"""
|
||||
client = _make_read_client(populated_storage)
|
||||
count = client.count_active_children("coord-1")
|
||||
# child-a (idle) + child-b (running) + child-coord (running/default) = 3
|
||||
assert count == 3
|
||||
|
||||
|
||||
def test_count_active_children_excludes_closed_and_deleted(populated_storage):
|
||||
"""A closed tail must not count toward the active-children budget —
|
||||
this is the whole reason for switching off list_children's
|
||||
LIMIT-then-filter path.
|
||||
"""
|
||||
# Close child-a and mark child-b deleted. child-coord stays active.
|
||||
populated_storage.update_workstream_state("child-a", "closed")
|
||||
populated_storage.update_workstream_state("child-b", "deleted")
|
||||
client = _make_read_client(populated_storage)
|
||||
count = client.count_active_children("coord-1")
|
||||
assert count == 1 # only child-coord survives
|
||||
|
||||
|
||||
def test_count_active_children_rejects_foreign_parent(populated_storage):
|
||||
"""Tenant guard — a crafted parent_ws_id other than the coord's own
|
||||
returns 0 without hitting storage."""
|
||||
client = _make_read_client(populated_storage)
|
||||
# The client's coord_ws_id is "coord-1" (see _make_read_client).
|
||||
# Counting against a different id must not leak anyone else's count.
|
||||
assert client.count_active_children("other-coord") == 0
|
||||
|
||||
|
||||
def test_count_active_children_fails_open_on_storage_error(populated_storage, monkeypatch):
|
||||
"""Budget is operator safety, not a security gate — a broken storage
|
||||
path must return 0 so the coord still makes progress."""
|
||||
client = _make_read_client(populated_storage)
|
||||
|
||||
def _boom(**_kwargs):
|
||||
raise RuntimeError("storage broken")
|
||||
|
||||
monkeypatch.setattr(populated_storage, "count_workstreams_by_state", _boom)
|
||||
assert client.count_active_children("coord-1") == 0
|
||||
|
||||
|
||||
def test_list_children_signals_truncation_when_page_full_and_filter_drops(
|
||||
populated_storage,
|
||||
):
|
||||
|
||||
@@ -1,395 +0,0 @@
|
||||
"""Tests for the coordinator ``/quota`` GET + POST endpoints.
|
||||
|
||||
Covers the admin partial-update surface for spawn-budget and
|
||||
spawn-rate — parallel to the /trust + /restrict shape in
|
||||
``test_coordinator_governance.py``. Kept in its own file so PR B's
|
||||
review surface stays tight.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.routing import Route
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from tests._coord_test_helpers import (
|
||||
_AuthMiddleware,
|
||||
_build_mgr,
|
||||
_fake_registry,
|
||||
_FakeConfigStore,
|
||||
)
|
||||
from turnstone.console.server import (
|
||||
coordinator_quota_get,
|
||||
coordinator_quota_post,
|
||||
)
|
||||
from turnstone.core.spawn_quota import SpawnBudget, TokenBucket
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path):
|
||||
return SQLiteBackend(str(tmp_path / "coord.db"))
|
||||
|
||||
|
||||
_COORD_HEADERS = {"X-Test-User": "user-1", "X-Test-Perms": "admin.coordinator"}
|
||||
|
||||
|
||||
def _make_client(storage, *, coord_mgr, alias="my-model", registry=None) -> TestClient:
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Route(
|
||||
"/v1/api/coordinator/{ws_id}/quota",
|
||||
coordinator_quota_get,
|
||||
methods=["GET"],
|
||||
),
|
||||
Route(
|
||||
"/v1/api/coordinator/{ws_id}/quota",
|
||||
coordinator_quota_post,
|
||||
methods=["POST"],
|
||||
),
|
||||
],
|
||||
middleware=[Middleware(_AuthMiddleware)],
|
||||
)
|
||||
app.state.coord_mgr = coord_mgr
|
||||
app.state.config_store = _FakeConfigStore({"coordinator.model_alias": alias})
|
||||
app.state.coord_registry = registry
|
||||
app.state.coord_registry_error = "" if coord_mgr else "registry missing"
|
||||
app.state.auth_storage = storage
|
||||
app.state.jwt_secret = "x" * 64
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def _install_quota(coord) -> tuple[SpawnBudget, TokenBucket]:
|
||||
"""Attach a real budget + bucket to the coord session under test."""
|
||||
budget = SpawnBudget(20)
|
||||
bucket = TokenBucket(5.0, 10)
|
||||
session = MagicMock()
|
||||
session._spawn_budget = budget
|
||||
session._spawn_bucket = bucket
|
||||
session._coord_client = MagicMock()
|
||||
|
||||
def _get_state():
|
||||
return {
|
||||
"spawn_budget": budget.budget,
|
||||
"spawn_rate": {
|
||||
"tokens_per_minute": bucket.tokens_per_minute,
|
||||
"burst": bucket.burst,
|
||||
"tokens_available": bucket.tokens,
|
||||
},
|
||||
}
|
||||
|
||||
def _set_budget(n):
|
||||
budget.set_budget(int(n))
|
||||
|
||||
def _set_rate(tpm, brst):
|
||||
bucket.set_rate(float(tpm), int(brst))
|
||||
|
||||
session.get_quota_state.side_effect = _get_state
|
||||
session.set_spawn_budget.side_effect = _set_budget
|
||||
session.set_spawn_rate.side_effect = _set_rate
|
||||
coord.session = session
|
||||
return budget, bucket
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_quota_get_returns_live_snapshot(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.get(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["status"] == "ok"
|
||||
assert body["spawn_budget"] == 20
|
||||
assert body["spawn_rate"]["tokens_per_minute"] == 5.0
|
||||
assert body["spawn_rate"]["burst"] == 10
|
||||
assert 0 <= body["spawn_rate"]["tokens_available"] <= 10
|
||||
|
||||
|
||||
def test_quota_get_404_when_session_not_loaded(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
coord.session = None
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.get(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# POST — happy path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_quota_post_updates_budget_only_and_audits(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
budget, bucket = _install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"spawn_budget": 42},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["spawn_budget"] == 42
|
||||
# Rate left untouched — the partial update didn't widen it.
|
||||
assert body["spawn_rate"]["tokens_per_minute"] == 5.0
|
||||
assert body["spawn_rate"]["burst"] == 10
|
||||
assert budget.budget == 42
|
||||
|
||||
events = [e for e in storage.list_audit_events() if e["action"] == "coordinator.quota.updated"]
|
||||
assert len(events) == 1
|
||||
detail = json.loads(events[0]["detail"])
|
||||
assert detail["before"]["spawn_budget"] == 20
|
||||
assert detail["after"]["spawn_budget"] == 42
|
||||
|
||||
|
||||
def test_quota_post_accepts_nested_spawn_rate(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_budget, bucket = _install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"spawn_rate": {"tokens_per_minute": 30.0, "burst": 15}},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["spawn_rate"]["tokens_per_minute"] == 30.0
|
||||
assert body["spawn_rate"]["burst"] == 15
|
||||
assert bucket.burst == 15
|
||||
|
||||
|
||||
def test_quota_post_accepts_flat_aliases(storage):
|
||||
"""The admin UI may flatten the rate object — both shapes must work."""
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_budget, bucket = _install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"tokens_per_minute": 12.0, "burst": 4},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert bucket.tokens_per_minute == 12.0
|
||||
assert bucket.burst == 4
|
||||
|
||||
|
||||
def test_quota_post_burst_only_preserves_refill_rate(storage):
|
||||
"""Changing only burst shouldn't zero the refill rate — a previous
|
||||
bug-prone shape in partial-update handlers that overwrite missing
|
||||
fields with defaults. Here the handler must read current state
|
||||
for the missing dimension."""
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_budget, bucket = _install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"burst": 3},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert bucket.tokens_per_minute == 5.0 # unchanged
|
||||
assert bucket.burst == 3
|
||||
|
||||
|
||||
def test_quota_post_updates_all_three_knobs_at_once(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
budget, bucket = _install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"spawn_budget": 50, "tokens_per_minute": 0.0, "burst": 1},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert budget.budget == 50
|
||||
assert bucket.tokens_per_minute == 0.0
|
||||
assert bucket.burst == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# POST — validation failures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_quota_post_rejects_empty_body(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_quota_post_rejects_out_of_range_budget(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
for bad in (0, -5, 10_000):
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"spawn_budget": bad},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 400, f"expected 400 for {bad}"
|
||||
|
||||
|
||||
def test_quota_post_rejects_non_numeric_rate(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"tokens_per_minute": "fast"},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_quota_post_rejects_out_of_range_rate(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
for bad_tpm in (-1.0, 1_000.0):
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"tokens_per_minute": bad_tpm},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_quota_post_rejects_out_of_range_burst(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
for bad in (0, -1, 10_000):
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"burst": bad},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_quota_post_rejects_mixed_nested_and_flat_body(storage):
|
||||
"""Schema description says 'don't mix' — the handler enforces it with 400.
|
||||
|
||||
Silently picking one side would make the admin UI's behaviour
|
||||
unpredictable when it accidentally sends both shapes (e.g. during
|
||||
a form-rewrite transition).
|
||||
"""
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"spawn_rate": {"burst": 5}, "burst": 9},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
assert "conflicting" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_quota_post_rejects_non_object_spawn_rate(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"spawn_rate": "not-an-object"},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_quota_post_rejects_bool_as_numeric_field(storage):
|
||||
"""``True`` passes ``isinstance(x, int)`` in Python — explicit reject."""
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
for payload in (
|
||||
{"spawn_budget": True},
|
||||
{"burst": True},
|
||||
{"tokens_per_minute": True},
|
||||
):
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json=payload,
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 400, f"expected 400 for {payload}"
|
||||
|
||||
|
||||
def test_quota_post_404_when_session_not_loaded(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
coord.session = None
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"spawn_budget": 5},
|
||||
headers=_COORD_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
def test_quota_post_without_admin_coordinator_is_rejected(storage):
|
||||
mgr = _build_mgr(storage)
|
||||
coord = mgr.create(user_id="user-1", name="coord-a")
|
||||
_install_quota(coord)
|
||||
|
||||
client = _make_client(storage, coord_mgr=mgr, registry=_fake_registry())
|
||||
resp = client.post(
|
||||
f"/v1/api/coordinator/{coord.id}/quota",
|
||||
json={"spawn_budget": 5},
|
||||
headers={"X-Test-User": "user-1", "X-Test-Perms": ""},
|
||||
)
|
||||
assert resp.status_code in (401, 403)
|
||||
@@ -1254,221 +1254,3 @@ def test_close_all_children_prepare_errors_when_coord_client_unavailable(coord_s
|
||||
item = sess._prepare_tool(_tc("close_all_children", {}))
|
||||
assert "error" in item
|
||||
assert "unavailable" in item["error"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Spawn quota (PR B) — budget + rate gates
|
||||
#
|
||||
# Gate fires at _prepare time so the model gets feedback without the
|
||||
# operator having to deny an approval. Budget is advisory (counts
|
||||
# active children from storage); rate consumes a token per ask.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_spawn_prepare_rejects_when_budget_reached(coord_session):
|
||||
sess, coord, _ui = coord_session
|
||||
sess.set_spawn_budget(3)
|
||||
coord.count_active_children.return_value = 3
|
||||
item = sess._prepare_tool(_tc("spawn_workstream", {"initial_message": "hi"}))
|
||||
assert "error" in item
|
||||
assert "spawn budget reached" in item["error"]
|
||||
assert "3/3" in item["error"]
|
||||
coord.spawn.assert_not_called()
|
||||
|
||||
|
||||
def test_spawn_prepare_allowed_when_below_budget(coord_session):
|
||||
sess, coord, _ui = coord_session
|
||||
sess.set_spawn_budget(5)
|
||||
coord.count_active_children.return_value = 2
|
||||
item = sess._prepare_tool(_tc("spawn_workstream", {"initial_message": "hi"}))
|
||||
assert "error" not in item
|
||||
assert item["needs_approval"] is True
|
||||
|
||||
|
||||
def test_spawn_prepare_rejects_when_rate_limited(coord_session):
|
||||
sess, _coord, _ui = coord_session
|
||||
# Zero refill + burst 1 → first call consumes; second fails with inf retry.
|
||||
sess.set_spawn_rate(tokens_per_minute=0.0, burst=1)
|
||||
sess._prepare_tool(_tc("spawn_workstream", {"initial_message": "first"}))
|
||||
item = sess._prepare_tool(
|
||||
_tc("spawn_workstream", {"initial_message": "second"}, call_id="call-2")
|
||||
)
|
||||
assert "error" in item
|
||||
assert "rate limited" in item["error"]
|
||||
|
||||
|
||||
def test_spawn_prepare_rate_limited_error_reports_retry_after(coord_session):
|
||||
sess, _coord, _ui = coord_session
|
||||
# 60 tokens/minute = 1/sec; drain, next acquire reports ~1s retry.
|
||||
sess.set_spawn_rate(tokens_per_minute=60.0, burst=1)
|
||||
sess._prepare_tool(_tc("spawn_workstream", {"initial_message": "first"}))
|
||||
item = sess._prepare_tool(
|
||||
_tc("spawn_workstream", {"initial_message": "second"}, call_id="call-2")
|
||||
)
|
||||
assert "error" in item
|
||||
assert "retry after" in item["error"]
|
||||
|
||||
|
||||
def test_spawn_batch_partial_success_when_budget_overflows(coord_session):
|
||||
"""Budget 4, active 3, batch of 3: first 1 passes, remaining 2 → denied."""
|
||||
sess, coord, _ui = coord_session
|
||||
sess.set_spawn_budget(4)
|
||||
coord.count_active_children.return_value = 3
|
||||
|
||||
counter = {"n": 0}
|
||||
|
||||
def _spawn(**_kwargs):
|
||||
counter["n"] += 1
|
||||
return {
|
||||
"ws_id": f"child-{counter['n']}",
|
||||
"name": "n",
|
||||
"node_id": "node",
|
||||
"status": 200,
|
||||
}
|
||||
|
||||
coord.spawn.side_effect = _spawn
|
||||
item = sess._prepare_tool(
|
||||
_tc(
|
||||
"spawn_batch",
|
||||
{
|
||||
"children": [
|
||||
{"initial_message": "A"},
|
||||
{"initial_message": "B"},
|
||||
{"initial_message": "C"},
|
||||
]
|
||||
},
|
||||
)
|
||||
)
|
||||
# The batch isn't rejected outright — one child still fits.
|
||||
assert "error" not in item
|
||||
_call_id, output = sess._exec_spawn_batch(item)
|
||||
body = json.loads(output)
|
||||
assert len(body["results"]) == 1
|
||||
assert len(body["denied"]) == 2
|
||||
assert all("spawn budget reached" in d["reason"] for d in body["denied"])
|
||||
# Only one spawn() call — the gate kept the other two from dispatching.
|
||||
assert counter["n"] == 1
|
||||
|
||||
|
||||
def test_spawn_batch_all_denied_by_budget_returns_tool_error(coord_session):
|
||||
"""Budget fully consumed → batch rejected without approval flow."""
|
||||
sess, coord, _ui = coord_session
|
||||
sess.set_spawn_budget(2)
|
||||
coord.count_active_children.return_value = 2
|
||||
item = sess._prepare_tool(
|
||||
_tc(
|
||||
"spawn_batch",
|
||||
{
|
||||
"children": [
|
||||
{"initial_message": "A"},
|
||||
{"initial_message": "B"},
|
||||
]
|
||||
},
|
||||
)
|
||||
)
|
||||
assert "error" in item
|
||||
assert "spawn budget reached" in item["error"]
|
||||
coord.spawn.assert_not_called()
|
||||
|
||||
|
||||
def test_spawn_batch_partial_success_when_rate_bucket_empties(coord_session):
|
||||
"""Burst 2, batch of 3: first 2 pass, third hits empty bucket → denied."""
|
||||
sess, coord, _ui = coord_session
|
||||
sess.set_spawn_rate(tokens_per_minute=0.0, burst=2)
|
||||
|
||||
counter = {"n": 0}
|
||||
|
||||
def _spawn(**_kwargs):
|
||||
counter["n"] += 1
|
||||
return {
|
||||
"ws_id": f"child-{counter['n']}",
|
||||
"name": "n",
|
||||
"node_id": "node",
|
||||
"status": 200,
|
||||
}
|
||||
|
||||
coord.spawn.side_effect = _spawn
|
||||
item = sess._prepare_tool(
|
||||
_tc(
|
||||
"spawn_batch",
|
||||
{
|
||||
"children": [
|
||||
{"initial_message": "A"},
|
||||
{"initial_message": "B"},
|
||||
{"initial_message": "C"},
|
||||
]
|
||||
},
|
||||
)
|
||||
)
|
||||
assert "error" not in item
|
||||
_call_id, output = sess._exec_spawn_batch(item)
|
||||
body = json.loads(output)
|
||||
assert len(body["results"]) == 2
|
||||
assert len(body["denied"]) == 1
|
||||
assert "rate limited" in body["denied"][0]["reason"]
|
||||
assert counter["n"] == 2
|
||||
|
||||
|
||||
def test_set_spawn_budget_mutates_live_session(coord_session):
|
||||
sess, coord, _ui = coord_session
|
||||
sess.set_spawn_budget(0) # no room for anything
|
||||
coord.count_active_children.return_value = 0
|
||||
item = sess._prepare_tool(_tc("spawn_workstream", {"initial_message": "hi"}))
|
||||
assert "error" in item
|
||||
assert "spawn budget reached" in item["error"]
|
||||
|
||||
|
||||
def test_get_quota_state_snapshot_shape(coord_session):
|
||||
sess, _coord, _ui = coord_session
|
||||
sess.set_spawn_budget(7)
|
||||
sess.set_spawn_rate(tokens_per_minute=12.0, burst=4)
|
||||
state = sess.get_quota_state()
|
||||
assert state["spawn_budget"] == 7
|
||||
assert state["spawn_rate"]["tokens_per_minute"] == 12.0
|
||||
assert state["spawn_rate"]["burst"] == 4
|
||||
# tokens_available reflects the live bucket state — should be <= burst.
|
||||
assert 0 <= state["spawn_rate"]["tokens_available"] <= 4
|
||||
|
||||
|
||||
def test_count_active_children_fails_open_on_storage_error(coord_session):
|
||||
"""The helper must fail *open* (return 0) so a broken storage path
|
||||
doesn't pin the coord to zero spawns — the budget is operator
|
||||
safety, not a security gate.
|
||||
"""
|
||||
sess, coord, _ui = coord_session
|
||||
coord.count_active_children.side_effect = RuntimeError("storage down")
|
||||
assert sess._count_active_children() == 0
|
||||
|
||||
|
||||
def test_non_coordinator_session_has_no_quota_state(monkeypatch):
|
||||
"""Interactive sessions shouldn't carry budget/bucket state — the
|
||||
quota is a coordinator-only surface. Exercises the real __init__
|
||||
path so a regression that flipped quota state on for non-coord
|
||||
sessions would actually break this test."""
|
||||
monkeypatch.setattr(ChatSession, "_load_skills", lambda self: None)
|
||||
monkeypatch.setattr(ChatSession, "_init_system_messages", lambda self: None)
|
||||
monkeypatch.setattr(ChatSession, "_save_config", lambda self: None)
|
||||
|
||||
sess = ChatSession(
|
||||
client=MagicMock(),
|
||||
model="gpt-test",
|
||||
ui=_StubUI(), # type: ignore[arg-type]
|
||||
instructions=None,
|
||||
temperature=0.0,
|
||||
max_tokens=1024,
|
||||
tool_timeout=30,
|
||||
context_window=16384,
|
||||
ws_id="interactive-1",
|
||||
user_id="user-1",
|
||||
client_type=ClientType.WEB,
|
||||
# kind defaults to INTERACTIVE; no coord_client.
|
||||
)
|
||||
assert sess._spawn_budget is None
|
||||
assert sess._spawn_bucket is None
|
||||
# The mutator helpers must no-op instead of crashing — a caller
|
||||
# that wandered into the quota surface on an interactive session
|
||||
# should get "nothing happened", not AttributeError.
|
||||
sess.set_spawn_budget(5)
|
||||
sess.set_spawn_rate(10.0, 3)
|
||||
assert sess._spawn_budget is None
|
||||
assert sess._spawn_bucket is None
|
||||
|
||||
@@ -1,142 +0,0 @@
|
||||
"""Unit tests for SpawnBudget + TokenBucket (turnstone/core/spawn_quota.py)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import time
|
||||
|
||||
from turnstone.core.spawn_quota import SpawnBudget, TokenBucket
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SpawnBudget
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_budget_below_cap_allows_spawn():
|
||||
b = SpawnBudget(5)
|
||||
res = b.check(active=2)
|
||||
assert res.allowed is True
|
||||
assert res.budget == 5
|
||||
assert res.active == 2
|
||||
assert res.remaining == 3
|
||||
|
||||
|
||||
def test_budget_at_cap_rejects_spawn():
|
||||
b = SpawnBudget(3)
|
||||
res = b.check(active=3)
|
||||
assert res.allowed is False
|
||||
assert res.remaining == 0
|
||||
|
||||
|
||||
def test_budget_over_cap_reports_zero_remaining():
|
||||
"""A stale active count above the cap still clamps remaining to 0."""
|
||||
b = SpawnBudget(3)
|
||||
res = b.check(active=5)
|
||||
assert res.allowed is False
|
||||
assert res.remaining == 0
|
||||
assert res.active == 5
|
||||
|
||||
|
||||
def test_budget_negative_active_normalised():
|
||||
"""A negative active value (shouldn't happen in practice) normalises to 0."""
|
||||
b = SpawnBudget(5)
|
||||
res = b.check(active=-3)
|
||||
assert res.allowed is True
|
||||
assert res.active == 0
|
||||
assert res.remaining == 5
|
||||
|
||||
|
||||
def test_budget_set_mutates_cap_live():
|
||||
b = SpawnBudget(5)
|
||||
b.set_budget(10)
|
||||
assert b.budget == 10
|
||||
assert b.check(active=7).allowed is True
|
||||
|
||||
|
||||
def test_budget_negative_constructor_clamps_to_zero():
|
||||
"""A defensive floor — budget=-1 shouldn't mean "infinite spawns"."""
|
||||
b = SpawnBudget(-5)
|
||||
assert b.budget == 0
|
||||
assert b.check(active=0).allowed is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TokenBucket
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_bucket_starts_full():
|
||||
"""Fresh buckets grant ``burst`` immediately — the rate limit is for
|
||||
pacing a runaway, not the first wave."""
|
||||
tb = TokenBucket(tokens_per_minute=6.0, burst=5)
|
||||
# Five acquires in a row succeed.
|
||||
for _ in range(5):
|
||||
assert tb.acquire().allowed is True
|
||||
# Sixth exhausts the bucket.
|
||||
ack = tb.acquire()
|
||||
assert ack.allowed is False
|
||||
assert ack.retry_after_seconds > 0.0
|
||||
|
||||
|
||||
def test_bucket_empty_reports_retry_after():
|
||||
tb = TokenBucket(tokens_per_minute=60.0, burst=1) # 1 token/sec
|
||||
tb.acquire() # drains
|
||||
ack = tb.acquire()
|
||||
assert ack.allowed is False
|
||||
# 1 token/sec → deficit 1.0 → retry ~1.0s
|
||||
assert math.isclose(ack.retry_after_seconds, 1.0, rel_tol=0.2)
|
||||
|
||||
|
||||
def test_bucket_zero_rate_reports_infinite_retry():
|
||||
"""A disabled rate (tokens_per_minute=0) shouldn't promise a retry."""
|
||||
tb = TokenBucket(tokens_per_minute=0.0, burst=2)
|
||||
tb.acquire()
|
||||
tb.acquire()
|
||||
ack = tb.acquire()
|
||||
assert ack.allowed is False
|
||||
assert ack.retry_after_seconds == float("inf")
|
||||
|
||||
|
||||
def test_bucket_refills_over_time():
|
||||
tb = TokenBucket(tokens_per_minute=600.0, burst=1) # 10 tokens/sec
|
||||
tb.acquire() # empty
|
||||
assert tb.acquire().allowed is False
|
||||
time.sleep(0.15) # ~1.5 tokens refilled; clamps to burst=1
|
||||
ack = tb.acquire()
|
||||
assert ack.allowed is True
|
||||
|
||||
|
||||
def test_bucket_refill_clamps_to_burst():
|
||||
tb = TokenBucket(tokens_per_minute=6000.0, burst=3) # 100/sec — saturates fast
|
||||
time.sleep(0.05) # easily enough to refill past burst
|
||||
for _ in range(3):
|
||||
assert tb.acquire().allowed is True
|
||||
# Fourth acquire must fail even after the long idle — burst caps retention.
|
||||
assert tb.acquire().allowed is False
|
||||
|
||||
|
||||
def test_bucket_set_rate_narrows_burst_immediately():
|
||||
tb = TokenBucket(tokens_per_minute=6.0, burst=10) # starts with 10 tokens
|
||||
tb.set_rate(tokens_per_minute=6.0, burst=3) # clamp down
|
||||
# Three succeed then exhausted.
|
||||
for _ in range(3):
|
||||
assert tb.acquire().allowed is True
|
||||
assert tb.acquire().allowed is False
|
||||
|
||||
|
||||
def test_bucket_set_rate_widening_does_not_grant_free_tokens():
|
||||
"""A widened burst shouldn't retroactively fill the bucket — operators
|
||||
adjusting quotas shouldn't accidentally green-light a burst."""
|
||||
tb = TokenBucket(tokens_per_minute=0.0, burst=2)
|
||||
tb.acquire()
|
||||
tb.acquire() # bucket drained
|
||||
tb.set_rate(tokens_per_minute=0.0, burst=10) # widen
|
||||
ack = tb.acquire()
|
||||
assert ack.allowed is False # still empty
|
||||
|
||||
|
||||
def test_bucket_tokens_property_is_snapshot():
|
||||
tb = TokenBucket(tokens_per_minute=0.0, burst=5)
|
||||
assert tb.tokens == 5.0
|
||||
tb.acquire()
|
||||
assert tb.tokens == 4.0
|
||||
@@ -1231,124 +1231,6 @@ class CoordinatorCloseAllChildrenResponse(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
# Quota-range bounds pulled from the settings registry so the OpenAPI
|
||||
# schema, the handler validator, and the admin UI all advertise the
|
||||
# same limits — one source of truth, no drift when an operator bumps a
|
||||
# setting's max_value.
|
||||
def _quota_field_bounds(key: str) -> tuple[float, float]:
|
||||
from turnstone.core.settings_registry import SETTINGS
|
||||
|
||||
defn = SETTINGS[key]
|
||||
lo = 0.0 if defn.min_value is None else float(defn.min_value)
|
||||
hi = float("inf") if defn.max_value is None else float(defn.max_value)
|
||||
return lo, hi
|
||||
|
||||
|
||||
_BUDGET_LO, _BUDGET_HI = _quota_field_bounds("coordinator.spawn_budget")
|
||||
_TPM_LO, _TPM_HI = _quota_field_bounds("coordinator.spawn_rate.tokens_per_minute")
|
||||
_BURST_LO, _BURST_HI = _quota_field_bounds("coordinator.spawn_rate.burst")
|
||||
|
||||
|
||||
class CoordinatorSpawnRateInput(BaseModel):
|
||||
"""Nested rate-bucket sub-object of the quota REQUEST body.
|
||||
|
||||
Intentionally excludes the response-only ``tokens_available`` field
|
||||
so generated SDK input types don't imply clients can POST a
|
||||
live-bucket reading — the server ignores it on input.
|
||||
"""
|
||||
|
||||
tokens_per_minute: float | None = Field(
|
||||
default=None,
|
||||
ge=_TPM_LO,
|
||||
le=_TPM_HI,
|
||||
description="New refill rate in tokens per minute.",
|
||||
)
|
||||
burst: int | None = Field(
|
||||
default=None,
|
||||
ge=int(_BURST_LO),
|
||||
le=int(_BURST_HI),
|
||||
description="New burst ceiling.",
|
||||
)
|
||||
|
||||
|
||||
class CoordinatorSpawnRateState(BaseModel):
|
||||
"""Nested rate-bucket sub-object of the quota RESPONSE body.
|
||||
|
||||
Extends the input shape with the read-only ``tokens_available``
|
||||
snapshot so the admin UI can render a "rate status" badge.
|
||||
"""
|
||||
|
||||
tokens_per_minute: float | None = Field(
|
||||
default=None,
|
||||
description="Current refill rate (tokens/minute). None on a non-coordinator session.",
|
||||
)
|
||||
burst: int | None = Field(
|
||||
default=None,
|
||||
description="Current burst ceiling. None on a non-coordinator session.",
|
||||
)
|
||||
tokens_available: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Current post-refill token count — advisory only; the bucket "
|
||||
"refills continuously so this snapshot is stale the moment "
|
||||
"it's read. Useful for the admin UI's 'rate status' badge. "
|
||||
"Response-only: the server ignores this field on input."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class CoordinatorQuotaRequest(BaseModel):
|
||||
"""Body for POST /v1/api/coordinator/{ws_id}/quota.
|
||||
|
||||
Partial-update semantics — any subset of the three knobs may be
|
||||
supplied; missing fields keep their current values. Accepts either
|
||||
the nested ``spawn_rate`` object OR the flat ``tokens_per_minute``
|
||||
/ ``burst`` aliases — supplying both for the same field yields a
|
||||
400 so the admin UI can't half-migrate its body shape unnoticed.
|
||||
"""
|
||||
|
||||
spawn_budget: int | None = Field(
|
||||
default=None,
|
||||
ge=int(_BUDGET_LO),
|
||||
le=int(_BUDGET_HI),
|
||||
description=f"New active-children cap ({int(_BUDGET_LO)}..{int(_BUDGET_HI)}).",
|
||||
)
|
||||
spawn_rate: CoordinatorSpawnRateInput | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Nested rate-bucket overrides. Use this OR the flat "
|
||||
"tokens_per_minute/burst aliases for a given field; mixing "
|
||||
"both shapes for the same field is rejected with 400."
|
||||
),
|
||||
)
|
||||
tokens_per_minute: float | None = Field(
|
||||
default=None,
|
||||
ge=_TPM_LO,
|
||||
le=_TPM_HI,
|
||||
description="Flat alias for spawn_rate.tokens_per_minute.",
|
||||
)
|
||||
burst: int | None = Field(
|
||||
default=None,
|
||||
ge=int(_BURST_LO),
|
||||
le=int(_BURST_HI),
|
||||
description="Flat alias for spawn_rate.burst.",
|
||||
)
|
||||
|
||||
|
||||
class CoordinatorQuotaResponse(BaseModel):
|
||||
"""Response body for GET/POST /v1/api/coordinator/{ws_id}/quota."""
|
||||
|
||||
status: str = Field(default="ok")
|
||||
spawn_budget: int | None = Field(
|
||||
default=None,
|
||||
description="Active-children cap. None on a non-coordinator session.",
|
||||
)
|
||||
spawn_rate: CoordinatorSpawnRateState = Field(
|
||||
default_factory=CoordinatorSpawnRateState,
|
||||
description="Rate-bucket state. Fields are None on non-coordinator sessions.",
|
||||
)
|
||||
|
||||
|
||||
class ClusterWsDetailResponse(BaseModel):
|
||||
"""Response body for GET /v1/api/cluster/ws/{ws_id}/detail.
|
||||
|
||||
|
||||
@@ -34,13 +34,9 @@ from turnstone.api.console_schemas import (
|
||||
CoordinatorInfo,
|
||||
CoordinatorListResponse,
|
||||
CoordinatorOpenResponse,
|
||||
CoordinatorQuotaRequest,
|
||||
CoordinatorQuotaResponse,
|
||||
CoordinatorRestrictRequest,
|
||||
CoordinatorRestrictResponse,
|
||||
CoordinatorSendRequest,
|
||||
CoordinatorSpawnRateInput,
|
||||
CoordinatorSpawnRateState,
|
||||
CoordinatorStopCascadeResponse,
|
||||
CoordinatorTaskInfo,
|
||||
CoordinatorTasksResponse,
|
||||
@@ -1386,40 +1382,6 @@ CONSOLE_ENDPOINTS: list[EndpointSpec] = [
|
||||
error_codes=[400, 403, 404, 503],
|
||||
tags=["Coordinator"],
|
||||
),
|
||||
EndpointSpec(
|
||||
"/v1/api/coordinator/{ws_id}/quota",
|
||||
"GET",
|
||||
"Read the coordinator session's live spawn-quota state",
|
||||
description=(
|
||||
"Returns the current ``spawn_budget`` (active-children cap) and "
|
||||
"``spawn_rate`` bucket (``tokens_per_minute``, ``burst``, "
|
||||
"``tokens_available``). Values reflect the in-memory override "
|
||||
"when an admin has mutated the session via POST; otherwise they "
|
||||
"reflect the global defaults baked in at session construction."
|
||||
),
|
||||
response_model=CoordinatorQuotaResponse,
|
||||
error_codes=[400, 403, 404, 503],
|
||||
tags=["Coordinator"],
|
||||
),
|
||||
EndpointSpec(
|
||||
"/v1/api/coordinator/{ws_id}/quota",
|
||||
"POST",
|
||||
"Mutate the coordinator session's spawn quota (partial update)",
|
||||
description=(
|
||||
"Updates any subset of ``spawn_budget``, "
|
||||
"``spawn_rate.tokens_per_minute``, and ``spawn_rate.burst``. "
|
||||
"Nested ``spawn_rate`` and flat ``tokens_per_minute`` / ``burst`` "
|
||||
"aliases are both accepted. Missing fields keep their current "
|
||||
"values. Overrides are in-memory only — a session reopen "
|
||||
"re-seeds from the global settings, matching the /trust and "
|
||||
"/restrict contract. Writes ``coordinator.quota.updated`` "
|
||||
"with the before/after snapshot."
|
||||
),
|
||||
request_model=CoordinatorQuotaRequest,
|
||||
response_model=CoordinatorQuotaResponse,
|
||||
error_codes=[400, 403, 404, 503],
|
||||
tags=["Coordinator"],
|
||||
),
|
||||
EndpointSpec(
|
||||
"/v1/api/cluster/ws/{ws_id}/detail",
|
||||
"GET",
|
||||
@@ -1495,13 +1457,9 @@ _ALL_MODELS: list[type[BaseModel]] = [
|
||||
CoordinatorInfo,
|
||||
CoordinatorListResponse,
|
||||
CoordinatorOpenResponse,
|
||||
CoordinatorQuotaRequest,
|
||||
CoordinatorQuotaResponse,
|
||||
CoordinatorRestrictRequest,
|
||||
CoordinatorRestrictResponse,
|
||||
CoordinatorSendRequest,
|
||||
CoordinatorSpawnRateInput,
|
||||
CoordinatorSpawnRateState,
|
||||
CoordinatorStopCascadeResponse,
|
||||
CoordinatorTaskInfo,
|
||||
CoordinatorTasksResponse,
|
||||
|
||||
@@ -853,44 +853,6 @@ class CoordinatorClient:
|
||||
truncated = len(raw) >= limit
|
||||
return {"children": children, "truncated": truncated}
|
||||
|
||||
# Terminal states that free a coordinator's spawn-budget slot — a
|
||||
# row in one of these has wound down (either soft-closed or
|
||||
# hard-deleted, whatever the enum the migration happens to use at
|
||||
# the time). Must match ``list_children``'s include_closed=False
|
||||
# filter so the budget and the UI's active-child view count the
|
||||
# same set.
|
||||
_BUDGET_TERMINAL_STATES: ClassVar[frozenset[str]] = frozenset({"closed", "deleted"})
|
||||
|
||||
def count_active_children(self, parent_ws_id: str) -> int:
|
||||
"""Return the count of non-terminal direct children.
|
||||
|
||||
Uses ``storage.count_workstreams_by_state`` (a SQL aggregate)
|
||||
rather than ``list_children``'s LIMIT-then-filter path — a
|
||||
fan-out with many recently-closed children would otherwise
|
||||
push live rows past the SQL LIMIT and silently undercount,
|
||||
letting spawn slots leak.
|
||||
|
||||
Tenant-guarded: a coord's LLM can't drive a count of someone
|
||||
else's subtree. Returns 0 on any storage error so the budget
|
||||
fails *open* rather than pinning the coord at zero spawns on a
|
||||
transient blip (operator-level safety, not a security gate).
|
||||
"""
|
||||
if parent_ws_id != self._coord_ws_id:
|
||||
return 0
|
||||
try:
|
||||
counts = self._storage.count_workstreams_by_state(
|
||||
parent_ws_id=parent_ws_id,
|
||||
user_id=self._user_id or None,
|
||||
)
|
||||
except Exception:
|
||||
log.debug(
|
||||
"coord_client.count_active_children.failed ws=%s",
|
||||
parent_ws_id[:8],
|
||||
exc_info=True,
|
||||
)
|
||||
return 0
|
||||
return sum(n for state, n in counts.items() if state not in self._BUDGET_TERMINAL_STATES)
|
||||
|
||||
# Auto-metadata keys that expose internal network topology (RFC 1918
|
||||
# addresses, interface maps) without contributing to any routing
|
||||
# decision a coordinator makes. Stripped from the default response
|
||||
|
||||
@@ -3701,181 +3701,6 @@ async def coordinator_close_all_children(request: Request) -> JSONResponse:
|
||||
)
|
||||
|
||||
|
||||
# Quota range bounds come from the settings registry so the admin UI,
|
||||
# Pydantic schema, and the endpoint stay in lockstep — bumping a cap
|
||||
# in one place shouldn't leave the others silently permissive.
|
||||
def _quota_range(key: str) -> tuple[float, float]:
|
||||
from turnstone.core.settings_registry import SETTINGS
|
||||
|
||||
defn = SETTINGS[key]
|
||||
lo = 0.0 if defn.min_value is None else float(defn.min_value)
|
||||
hi = float("inf") if defn.max_value is None else float(defn.max_value)
|
||||
return lo, hi
|
||||
|
||||
|
||||
_QUOTA_MIN_BUDGET, _QUOTA_MAX_BUDGET = _quota_range("coordinator.spawn_budget")
|
||||
_QUOTA_MIN_TPM, _QUOTA_MAX_TPM = _quota_range("coordinator.spawn_rate.tokens_per_minute")
|
||||
_QUOTA_MIN_BURST, _QUOTA_MAX_BURST = _quota_range("coordinator.spawn_rate.burst")
|
||||
|
||||
|
||||
async def coordinator_quota_get(request: Request) -> JSONResponse:
|
||||
"""GET /v1/api/coordinator/{ws_id}/quota — read live spawn-quota state.
|
||||
|
||||
``allow_service_bypass=False`` so a service token whose ``user_id``
|
||||
matches the coord owner still needs an explicit ``admin.coordinator``
|
||||
grant. Consistent with the POST below and with the other
|
||||
destructive / capability-visibility coordinator endpoints
|
||||
(``/restrict``, ``/stop_cascade``, ``/close_all_children``).
|
||||
"""
|
||||
resolved = await _resolve_coord_session(request, allow_service_bypass=False)
|
||||
if isinstance(resolved, JSONResponse):
|
||||
return resolved
|
||||
session, _storage, _user_id, _ws_id = resolved
|
||||
state = session.get_quota_state()
|
||||
return JSONResponse({"status": "ok", **state})
|
||||
|
||||
|
||||
async def coordinator_quota_post(request: Request) -> JSONResponse:
|
||||
"""POST /v1/api/coordinator/{ws_id}/quota — mutate the session's spawn quota.
|
||||
|
||||
Partial-update body: any subset of ``spawn_budget``,
|
||||
``spawn_rate.tokens_per_minute``, ``spawn_rate.burst`` (or the
|
||||
flat aliases ``tokens_per_minute`` / ``burst``). Missing fields
|
||||
keep their current values. Overrides are in-memory only — a
|
||||
session reopen re-seeds from the global settings, matching the
|
||||
/trust and /restrict contract.
|
||||
|
||||
Gated with ``allow_service_bypass=False`` because this endpoint
|
||||
can RAISE a coord's spawn capacity — a service token whose
|
||||
``user_id`` matches the coord owner still needs an explicit
|
||||
``admin.coordinator`` grant, matching how ``/restrict`` and the
|
||||
cascade endpoints treat destructive / capability-escalating
|
||||
surfaces.
|
||||
"""
|
||||
resolved = await _resolve_coord_session(request, allow_service_bypass=False)
|
||||
if isinstance(resolved, JSONResponse):
|
||||
return resolved
|
||||
session, storage, user_id, ws_id = resolved
|
||||
|
||||
body = await _require_json_object(request)
|
||||
if isinstance(body, JSONResponse):
|
||||
return body
|
||||
|
||||
before = session.get_quota_state()
|
||||
|
||||
raw_budget = body.get("spawn_budget", None)
|
||||
new_budget: int | None = None
|
||||
if raw_budget is not None:
|
||||
if isinstance(raw_budget, bool) or not isinstance(raw_budget, int):
|
||||
return JSONResponse({"error": "spawn_budget must be an integer"}, status_code=400)
|
||||
if raw_budget < _QUOTA_MIN_BUDGET or raw_budget > _QUOTA_MAX_BUDGET:
|
||||
return JSONResponse(
|
||||
{
|
||||
"error": (
|
||||
f"spawn_budget out of range ({int(_QUOTA_MIN_BUDGET)}.."
|
||||
f"{int(_QUOTA_MAX_BUDGET)})"
|
||||
)
|
||||
},
|
||||
status_code=400,
|
||||
)
|
||||
new_budget = raw_budget
|
||||
|
||||
# Accept either a nested ``spawn_rate`` object or flat aliases so
|
||||
# the admin UI doesn't have to know which shape the session expects.
|
||||
# Reject on conflict (both nested and flat present for the same
|
||||
# field) — the schema description says "don't mix" and silently
|
||||
# picking one would make the admin UI behaviour unpredictable.
|
||||
rate_src: dict[str, Any] = {}
|
||||
raw_rate = body.get("spawn_rate")
|
||||
if raw_rate is not None and not isinstance(raw_rate, dict):
|
||||
return JSONResponse(
|
||||
{"error": "spawn_rate must be an object"},
|
||||
status_code=400,
|
||||
)
|
||||
if isinstance(raw_rate, dict):
|
||||
rate_src.update(raw_rate)
|
||||
for flat in ("tokens_per_minute", "burst"):
|
||||
if flat in body:
|
||||
if flat in rate_src:
|
||||
return JSONResponse(
|
||||
{
|
||||
"error": (
|
||||
f"conflicting nested spawn_rate.{flat} and flat "
|
||||
f"{flat} in body; use one or the other"
|
||||
)
|
||||
},
|
||||
status_code=400,
|
||||
)
|
||||
rate_src[flat] = body[flat]
|
||||
|
||||
new_tpm: float | None = None
|
||||
new_burst: int | None = None
|
||||
if "tokens_per_minute" in rate_src:
|
||||
raw_tpm = rate_src["tokens_per_minute"]
|
||||
if isinstance(raw_tpm, bool) or not isinstance(raw_tpm, (int, float)):
|
||||
return JSONResponse({"error": "tokens_per_minute must be a number"}, status_code=400)
|
||||
tpm = float(raw_tpm)
|
||||
if tpm < _QUOTA_MIN_TPM or tpm > _QUOTA_MAX_TPM:
|
||||
return JSONResponse(
|
||||
{"error": (f"tokens_per_minute out of range ({_QUOTA_MIN_TPM}..{_QUOTA_MAX_TPM})")},
|
||||
status_code=400,
|
||||
)
|
||||
new_tpm = tpm
|
||||
if "burst" in rate_src:
|
||||
raw_burst = rate_src["burst"]
|
||||
if isinstance(raw_burst, bool) or not isinstance(raw_burst, int):
|
||||
return JSONResponse({"error": "burst must be an integer"}, status_code=400)
|
||||
if raw_burst < _QUOTA_MIN_BURST or raw_burst > _QUOTA_MAX_BURST:
|
||||
return JSONResponse(
|
||||
{
|
||||
"error": (
|
||||
f"burst out of range ({int(_QUOTA_MIN_BURST)}..{int(_QUOTA_MAX_BURST)})"
|
||||
)
|
||||
},
|
||||
status_code=400,
|
||||
)
|
||||
new_burst = raw_burst
|
||||
|
||||
if new_budget is None and new_tpm is None and new_burst is None:
|
||||
return JSONResponse(
|
||||
{
|
||||
"error": (
|
||||
"body must carry at least one of spawn_budget, "
|
||||
"spawn_rate.tokens_per_minute, spawn_rate.burst"
|
||||
)
|
||||
},
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
if new_budget is not None:
|
||||
session.set_spawn_budget(new_budget)
|
||||
if new_tpm is not None or new_burst is not None:
|
||||
# set_spawn_rate takes both args — fill missing ones from the
|
||||
# current live state so a "change only burst" call doesn't
|
||||
# silently reset the refill rate.
|
||||
cur_rate = before.get("spawn_rate") or {}
|
||||
tpm_val = (
|
||||
new_tpm if new_tpm is not None else float(cur_rate.get("tokens_per_minute") or 0.0)
|
||||
)
|
||||
burst_val = new_burst if new_burst is not None else int(cur_rate.get("burst") or 1)
|
||||
session.set_spawn_rate(tpm_val, burst_val)
|
||||
|
||||
after = session.get_quota_state()
|
||||
await _emit_coord_audit(
|
||||
storage,
|
||||
user_id,
|
||||
"coordinator.quota.updated",
|
||||
ws_id,
|
||||
{
|
||||
"src": "coordinator",
|
||||
"before": before,
|
||||
"after": after,
|
||||
},
|
||||
request.client.host if request.client else "",
|
||||
)
|
||||
return JSONResponse({"status": "ok", **after})
|
||||
|
||||
|
||||
async def coordinator_tasks(request: Request) -> JSONResponse:
|
||||
"""GET /v1/api/coordinator/{ws_id}/tasks — read task list envelope.
|
||||
|
||||
@@ -10557,16 +10382,6 @@ def create_app(
|
||||
coordinator_close_all_children,
|
||||
methods=["POST"],
|
||||
),
|
||||
Route(
|
||||
"/api/coordinator/{ws_id}/quota",
|
||||
coordinator_quota_get,
|
||||
methods=["GET"],
|
||||
),
|
||||
Route(
|
||||
"/api/coordinator/{ws_id}/quota",
|
||||
coordinator_quota_post,
|
||||
methods=["POST"],
|
||||
),
|
||||
Route(
|
||||
"/api/coordinator/{ws_id}",
|
||||
coordinator_detail,
|
||||
|
||||
+6
-160
@@ -89,7 +89,6 @@ from turnstone.core.metacognition import (
|
||||
from turnstone.core.providers import create_provider
|
||||
from turnstone.core.safety import is_command_blocked, sanitize_command
|
||||
from turnstone.core.sandbox import execute_math_sandboxed
|
||||
from turnstone.core.spawn_quota import SpawnBudget, TokenBucket
|
||||
from turnstone.core.storage._registry import get_storage
|
||||
from turnstone.core.tool_search import ToolSearchManager
|
||||
from turnstone.core.tools import (
|
||||
@@ -347,11 +346,6 @@ class ChatSession:
|
||||
self._trust_send: bool = False
|
||||
self._revoked_tools: frozenset[str] = frozenset()
|
||||
self._governance_lock = threading.Lock()
|
||||
# Spawn quota state — populated post-config-store init when
|
||||
# ``kind == COORDINATOR``. Non-coord sessions leave these at
|
||||
# None; the ``_prepare_spawn_*`` gates short-circuit on that.
|
||||
self._spawn_budget: SpawnBudget | None = None
|
||||
self._spawn_bucket: TokenBucket | None = None
|
||||
self._registry = registry
|
||||
self._model_alias = model_alias
|
||||
self._health_registry = health_registry
|
||||
@@ -387,35 +381,6 @@ class ChatSession:
|
||||
self._username = username
|
||||
self._client_type = client_type
|
||||
self._config_store = config_store
|
||||
# Coordinator-only: spawn quota state. Global defaults live in
|
||||
# ``settings_registry``; per-session overrides land via
|
||||
# /v1/api/coordinator/{ws_id}/quota (in-memory, dies on session
|
||||
# reopen — matches the /trust and /restrict contract).
|
||||
if kind == WorkstreamKind.COORDINATOR:
|
||||
budget_default = 20
|
||||
tpm_default = 5.0
|
||||
burst_default = 10
|
||||
if config_store is not None:
|
||||
try:
|
||||
# Explicit None-check instead of ``or fallback`` so a
|
||||
# legitimate ``tokens_per_minute=0`` (the "disable
|
||||
# refill" operator knob) doesn't silently become the
|
||||
# default. The same pattern on budget / burst is
|
||||
# safe today (min_value=1 in the registry) but kept
|
||||
# consistent against a future range relaxation.
|
||||
raw_budget = config_store.get("coordinator.spawn_budget")
|
||||
if raw_budget is not None:
|
||||
budget_default = int(raw_budget)
|
||||
raw_tpm = config_store.get("coordinator.spawn_rate.tokens_per_minute")
|
||||
if raw_tpm is not None:
|
||||
tpm_default = float(raw_tpm)
|
||||
raw_burst = config_store.get("coordinator.spawn_rate.burst")
|
||||
if raw_burst is not None:
|
||||
burst_default = int(raw_burst)
|
||||
except (TypeError, ValueError):
|
||||
log.debug("coord_quota.config_read_failed", exc_info=True)
|
||||
self._spawn_budget = SpawnBudget(budget_default)
|
||||
self._spawn_bucket = TokenBucket(tpm_default, burst_default)
|
||||
# Initialize rule registry for configurable judge rules
|
||||
self._rule_registry = None
|
||||
if config_store is not None:
|
||||
@@ -4920,100 +4885,6 @@ class ChatSession:
|
||||
def get_revoked_tools(self) -> frozenset[str]:
|
||||
return self._revoked_tools
|
||||
|
||||
# -- Coordinator spawn quota (budget + rate) -------------------------
|
||||
#
|
||||
# Admin endpoint POST /v1/api/coordinator/{ws_id}/quota mutates the
|
||||
# live session state via ``set_spawn_budget`` / ``set_spawn_rate``.
|
||||
# Overrides are in-memory only; a session reopen re-seeds from the
|
||||
# global settings (mirrors /trust + /restrict).
|
||||
|
||||
def set_spawn_budget(self, budget: int) -> None:
|
||||
if self._spawn_budget is not None:
|
||||
self._spawn_budget.set_budget(int(budget))
|
||||
|
||||
def set_spawn_rate(self, tokens_per_minute: float, burst: int) -> None:
|
||||
if self._spawn_bucket is not None:
|
||||
self._spawn_bucket.set_rate(float(tokens_per_minute), int(burst))
|
||||
|
||||
def get_quota_state(self) -> dict[str, Any]:
|
||||
"""Snapshot of the coord-session quota for /quota GET + audit."""
|
||||
state: dict[str, Any] = {
|
||||
"spawn_budget": None,
|
||||
"spawn_rate": {"tokens_per_minute": None, "burst": None, "tokens_available": None},
|
||||
}
|
||||
if self._spawn_budget is not None:
|
||||
state["spawn_budget"] = self._spawn_budget.budget
|
||||
if self._spawn_bucket is not None:
|
||||
state["spawn_rate"] = {
|
||||
"tokens_per_minute": self._spawn_bucket.tokens_per_minute,
|
||||
"burst": self._spawn_bucket.burst,
|
||||
"tokens_available": self._spawn_bucket.tokens,
|
||||
}
|
||||
return state
|
||||
|
||||
def _count_active_children(self) -> int:
|
||||
"""Current active-child count for the budget check.
|
||||
|
||||
Routes through ``CoordinatorClient.count_active_children`` —
|
||||
an aggregate SQL count excluding ``closed`` / ``deleted`` —
|
||||
so the budget stays accurate even when the coord has a long
|
||||
tail of closed rows that would otherwise push live rows past
|
||||
the ``list_children`` LIMIT and silently undercount. Returns
|
||||
0 on any lookup error so the budget fails *open* rather than
|
||||
pinning the coord at zero spawns on a transient storage blip
|
||||
(the budget is operator-level safety, not a security gate).
|
||||
"""
|
||||
if self._coord_client is None or self._spawn_budget is None:
|
||||
return 0
|
||||
try:
|
||||
return int(self._coord_client.count_active_children(self._ws_id))
|
||||
except Exception:
|
||||
log.debug("coord_quota.count_active_failed", exc_info=True)
|
||||
return 0
|
||||
|
||||
def _eval_spawn_quota(self, active: int) -> str | None:
|
||||
"""Core budget + rate gate — shared by single-spawn and batch paths.
|
||||
|
||||
Runs the budget check against ``active`` (caller's running tally
|
||||
including any in-batch approvals) then consumes a token from
|
||||
the rate bucket. Returns a flat denial reason on failure, or
|
||||
``None`` on pass. A pass consumes a rate token, so callers
|
||||
must NOT re-evaluate the same attempt. Single source of truth
|
||||
for the error wording so the single-spawn path and the batch
|
||||
``_error`` rows stay in lockstep.
|
||||
"""
|
||||
if self._spawn_budget is not None:
|
||||
res = self._spawn_budget.check(active)
|
||||
if not res.allowed:
|
||||
return (
|
||||
f"spawn budget reached ({res.active}/{res.budget}); "
|
||||
"close idle children with close_workstream(ws_id) "
|
||||
"before spawning more"
|
||||
)
|
||||
if self._spawn_bucket is not None:
|
||||
ack = self._spawn_bucket.acquire()
|
||||
if not ack.allowed:
|
||||
hint = (
|
||||
f"retry after {ack.retry_after_seconds:.1f}s"
|
||||
if ack.retry_after_seconds != float("inf")
|
||||
else "rate limit disabled (tokens_per_minute=0)"
|
||||
)
|
||||
return f"spawn rate limited; {hint}"
|
||||
return None
|
||||
|
||||
def _check_spawn_quota(self, call_id: str, func_name: str) -> dict[str, Any] | None:
|
||||
"""Single-spawn gate wrapper.
|
||||
|
||||
Counts active children, evaluates the quota, and wraps any
|
||||
denial reason in ``_coord_tool_error`` so the preparer can
|
||||
return it directly. A pass consumes a rate token, so callers
|
||||
must NOT re-check.
|
||||
"""
|
||||
reason = self._eval_spawn_quota(self._count_active_children())
|
||||
if reason is None:
|
||||
return None
|
||||
return self._coord_tool_error(call_id, func_name, reason)
|
||||
|
||||
@staticmethod
|
||||
def _coord_str_arg(args: dict[str, Any], key: str, default: str = "") -> str:
|
||||
"""Return ``args[key]`` if it's a string, else ``default``.
|
||||
@@ -5054,15 +4925,6 @@ class ChatSession:
|
||||
return self._coord_tool_error(
|
||||
call_id, "spawn_workstream", "coordinator client unavailable"
|
||||
)
|
||||
# Quota gate runs before approval so the model sees a budget /
|
||||
# rate error immediately — avoids the operator having to reject
|
||||
# an approval that would have been denied anyway. The bucket
|
||||
# token is consumed here; a later operator-deny does NOT refund
|
||||
# it (the model burnt an ask, the pacing limit should reflect
|
||||
# that).
|
||||
quota_err = self._check_spawn_quota(call_id, "spawn_workstream")
|
||||
if quota_err is not None:
|
||||
return quota_err
|
||||
# Empty initial_message is allowed — creates an idle child
|
||||
# workstream ready to receive the first turn via
|
||||
# send_to_workstream. The tool JSON advertises this explicitly.
|
||||
@@ -5196,28 +5058,12 @@ class ChatSession:
|
||||
else:
|
||||
preview_rows.append(f" {idx}. (idle)")
|
||||
|
||||
# Partial-success quota gate via the shared ``_eval_spawn_quota``
|
||||
# helper — single source of truth with ``_prepare_spawn_workstream``
|
||||
# for the error wording + consume-on-ask semantics. Denied rows
|
||||
# get an ``_error`` field so the exec loop routes them to
|
||||
# ``denied[]`` without re-checking. Tokens consumed here are
|
||||
# NOT refunded on operator-deny — matches single-spawn semantics.
|
||||
active_count = self._count_active_children()
|
||||
approved_in_batch = 0
|
||||
for spec in normalised:
|
||||
if "_error" in spec:
|
||||
continue
|
||||
reason = self._eval_spawn_quota(active_count + approved_in_batch)
|
||||
if reason is not None:
|
||||
spec["_error"] = reason
|
||||
continue
|
||||
approved_in_batch += 1
|
||||
|
||||
# If every row was denied by the quota gate (or invalid at
|
||||
# normalisation), skip the approval round — operators shouldn't
|
||||
# approve a batch with nothing to spawn. Surface the first
|
||||
# denial reason directly so the model gets actionable feedback.
|
||||
if approved_in_batch == 0:
|
||||
# If every row was invalid at normalisation, skip the approval
|
||||
# round — operators shouldn't approve a batch with nothing to
|
||||
# spawn. Surface the first denial reason directly so the model
|
||||
# gets actionable feedback.
|
||||
valid_count = sum(1 for spec in normalised if "_error" not in spec)
|
||||
if valid_count == 0:
|
||||
first_err = next(
|
||||
(s.get("_error") for s in normalised if s.get("_error")), "batch rejected"
|
||||
)
|
||||
|
||||
@@ -676,47 +676,6 @@ def _build_registry() -> dict[str, SettingDef]:
|
||||
"this short (5 minutes default) to limit blast radius if the process "
|
||||
"is compromised; longer values reduce JWT re-mint frequency at minor risk.",
|
||||
),
|
||||
SettingDef(
|
||||
"coordinator.spawn_budget",
|
||||
"int",
|
||||
20,
|
||||
"Max concurrent active children per coordinator (hard quota)",
|
||||
"coordinator",
|
||||
min_value=1,
|
||||
max_value=500,
|
||||
help="Hard cap on how many children a single coordinator session can "
|
||||
"own at once (counts non-terminal children only — closed / deleted rows "
|
||||
"free a slot). Once at the cap, `spawn_workstream` returns a tool error "
|
||||
"guiding the model to close idle children first. The per-session value "
|
||||
"can be overridden at runtime via POST /v1/api/coordinator/{ws_id}/quota.",
|
||||
),
|
||||
SettingDef(
|
||||
"coordinator.spawn_rate.tokens_per_minute",
|
||||
"float",
|
||||
5.0,
|
||||
"Token-bucket refill rate for coordinator spawns (tokens/minute)",
|
||||
"coordinator",
|
||||
min_value=0.0,
|
||||
max_value=600.0,
|
||||
help="Soft pacing limit — how fast a coordinator can sustain spawns once "
|
||||
"the initial burst is exhausted. A rate-limited spawn returns a tool error "
|
||||
"with a `retry_after_seconds` hint so the model backs off before trying "
|
||||
"again. Set to 0 to disable the refill (bucket still honors the initial "
|
||||
"burst on a fresh session).",
|
||||
),
|
||||
SettingDef(
|
||||
"coordinator.spawn_rate.burst",
|
||||
"int",
|
||||
10,
|
||||
"Token-bucket burst ceiling for coordinator spawns",
|
||||
"coordinator",
|
||||
min_value=1,
|
||||
max_value=500,
|
||||
help="Max tokens the spawn-rate bucket can accrue — the initial fan-out "
|
||||
"size a fresh or long-idle coordinator can issue back-to-back before the "
|
||||
"`tokens_per_minute` refill rate takes over. Larger values trade slower "
|
||||
"convergence to steady state for more elastic first-wave fan-outs.",
|
||||
),
|
||||
]
|
||||
return {d.key: d for d in defs}
|
||||
|
||||
|
||||
@@ -1,183 +0,0 @@
|
||||
"""Spawn-rate / spawn-budget controls for coordinator sessions.
|
||||
|
||||
Two complementary knobs layered on top of ``spawn_workstream`` /
|
||||
``spawn_batch``:
|
||||
|
||||
- :class:`SpawnBudget` — hard cap on the number of *concurrently
|
||||
active* children a coordinator can own. A cheap in-memory struct;
|
||||
the authoritative active count lives in storage and the caller
|
||||
passes it in via :meth:`SpawnBudget.check`. Keeping the count out
|
||||
of this class keeps the budget trivial to unit-test without a
|
||||
storage fixture.
|
||||
- :class:`TokenBucket` — soft pacing of spawn attempts. Refills at a
|
||||
configurable per-minute rate up to a burst ceiling. ``acquire()``
|
||||
either consumes a token or returns a ``retry_after_seconds`` hint
|
||||
the model can include in its next planning turn.
|
||||
|
||||
Both live on the coordinator's :class:`ChatSession` (one pair per
|
||||
session) and share a tiny ``threading.Lock`` each — the approval layer
|
||||
can interleave with SSE observers mid-prepare, so the mutators must
|
||||
not race with the reads.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BudgetCheck:
|
||||
"""Result of :meth:`SpawnBudget.check`.
|
||||
|
||||
``allowed`` is True when ``active + cost <= budget``; ``remaining``
|
||||
is clamped at zero so the caller can safely do
|
||||
``min(remaining, len(batch))`` to size a partial-success batch.
|
||||
"""
|
||||
|
||||
allowed: bool
|
||||
active: int
|
||||
budget: int
|
||||
remaining: int
|
||||
|
||||
|
||||
class SpawnBudget:
|
||||
"""Thread-safe hard cap on concurrent active children."""
|
||||
|
||||
def __init__(self, budget: int) -> None:
|
||||
self._budget = max(0, int(budget))
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def budget(self) -> int:
|
||||
with self._lock:
|
||||
return self._budget
|
||||
|
||||
def set_budget(self, budget: int) -> None:
|
||||
with self._lock:
|
||||
self._budget = max(0, int(budget))
|
||||
|
||||
def check(self, active: int) -> BudgetCheck:
|
||||
"""Return whether one more spawn fits given the current active count.
|
||||
|
||||
Does NOT mutate state — batch callers are expected to re-check
|
||||
per item with their running in-batch tally added to ``active``
|
||||
(a single multi-spawn ``cost`` kwarg would atomically fast-fail
|
||||
the whole tail; add it back if that semantic lands).
|
||||
"""
|
||||
with self._lock:
|
||||
b = self._budget
|
||||
clamped = max(0, active)
|
||||
remaining = max(0, b - clamped)
|
||||
return BudgetCheck(
|
||||
allowed=(clamped + 1) <= b,
|
||||
active=clamped,
|
||||
budget=b,
|
||||
remaining=remaining,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BucketAcquire:
|
||||
"""Result of :meth:`TokenBucket.acquire`.
|
||||
|
||||
On ``allowed=True`` the token has already been consumed; on
|
||||
``allowed=False`` ``retry_after_seconds`` is the wall-clock delay
|
||||
until enough tokens exist to satisfy the call (or ``inf`` when
|
||||
the refill rate is zero).
|
||||
"""
|
||||
|
||||
allowed: bool
|
||||
retry_after_seconds: float
|
||||
tokens_remaining: float
|
||||
|
||||
|
||||
class TokenBucket:
|
||||
"""Classic token bucket — tokens refill at ``tokens_per_minute`` up to ``burst``.
|
||||
|
||||
Starts full so a fresh session gets its full burst immediately;
|
||||
the rate limit exists to pace a *runaway*, not to speed-bump the
|
||||
first handful of spawns.
|
||||
"""
|
||||
|
||||
def __init__(self, tokens_per_minute: float, burst: int) -> None:
|
||||
self._per_sec = max(0.0, float(tokens_per_minute)) / 60.0
|
||||
self._burst = max(1, int(burst))
|
||||
self._tokens = float(self._burst)
|
||||
self._last_refill = time.monotonic()
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def tokens_per_minute(self) -> float:
|
||||
with self._lock:
|
||||
return self._per_sec * 60.0
|
||||
|
||||
@property
|
||||
def burst(self) -> int:
|
||||
with self._lock:
|
||||
return self._burst
|
||||
|
||||
@property
|
||||
def tokens(self) -> float:
|
||||
"""Current token count, post-refill. Intended for admin-read only."""
|
||||
now = time.monotonic()
|
||||
with self._lock:
|
||||
self._refill_locked(now)
|
||||
return self._tokens
|
||||
|
||||
def set_rate(self, tokens_per_minute: float, burst: int) -> None:
|
||||
"""Mutate the refill rate and/or burst ceiling.
|
||||
|
||||
The current token count is clamped DOWN to the new burst — a
|
||||
widened burst doesn't retroactively grant free tokens to a
|
||||
bucket that was sitting full. A narrowed burst takes effect
|
||||
immediately.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
with self._lock:
|
||||
self._refill_locked(now)
|
||||
self._per_sec = max(0.0, float(tokens_per_minute)) / 60.0
|
||||
self._burst = max(1, int(burst))
|
||||
if self._tokens > self._burst:
|
||||
self._tokens = float(self._burst)
|
||||
|
||||
def _refill_locked(self, now: float) -> None:
|
||||
elapsed = max(0.0, now - self._last_refill)
|
||||
self._tokens = min(float(self._burst), self._tokens + elapsed * self._per_sec)
|
||||
self._last_refill = now
|
||||
|
||||
def acquire(self, *, cost: float = 1.0) -> BucketAcquire:
|
||||
"""Try to consume ``cost`` tokens.
|
||||
|
||||
Returns ``(True, 0.0, remaining)`` on success; on failure
|
||||
returns ``(False, retry_after, current)`` where
|
||||
``retry_after`` is the time until enough tokens exist. When
|
||||
``tokens_per_minute == 0`` a failing acquire reports
|
||||
``retry_after == inf`` — the caller should surface that as
|
||||
"rate limit disabled; change the setting" rather than a
|
||||
real retry hint.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
with self._lock:
|
||||
self._refill_locked(now)
|
||||
if self._tokens >= cost:
|
||||
self._tokens -= cost
|
||||
return BucketAcquire(
|
||||
allowed=True,
|
||||
retry_after_seconds=0.0,
|
||||
tokens_remaining=self._tokens,
|
||||
)
|
||||
if self._per_sec <= 0:
|
||||
return BucketAcquire(
|
||||
allowed=False,
|
||||
retry_after_seconds=float("inf"),
|
||||
tokens_remaining=self._tokens,
|
||||
)
|
||||
deficit = cost - self._tokens
|
||||
retry_after = deficit / self._per_sec
|
||||
return BucketAcquire(
|
||||
allowed=False,
|
||||
retry_after_seconds=retry_after,
|
||||
tokens_remaining=self._tokens,
|
||||
)
|
||||
@@ -0,0 +1,44 @@
|
||||
"""Drop coord spawn-quota settings rows.
|
||||
|
||||
The coordinator spawn-quota subsystem (``SpawnBudget`` / ``TokenBucket``
|
||||
gating ``spawn_workstream``) is gone — ``max_active`` slot exhaustion
|
||||
already bounds runaway spawns and surfaces as a tool error. The three
|
||||
registry keys that configured it have been removed from
|
||||
``settings_registry``; this migration clears any persisted rows so
|
||||
upgraded deployments don't log "Skipping invalid setting" warnings
|
||||
for them on every startup.
|
||||
|
||||
Downgrade is a no-op: the deleted rows were operator-set values; once
|
||||
gone they're gone. If a deployment rolls back to pre-1.5.0 code, the
|
||||
registry at that point still has the keys and ``reload()`` falls back
|
||||
to the registry default (20 / 5.0 / 10) for any key not present in
|
||||
``system_settings``.
|
||||
|
||||
Revision ID: 047
|
||||
Revises: 046
|
||||
Create Date: 2026-04-23
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "047"
|
||||
down_revision = "046"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
bind.execute(
|
||||
sa.text(
|
||||
"DELETE FROM system_settings WHERE key IN ("
|
||||
"'coordinator.spawn_budget', "
|
||||
"'coordinator.spawn_rate.tokens_per_minute', "
|
||||
"'coordinator.spawn_rate.burst')"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
Reference in New Issue
Block a user