diff --git a/sdk/typescript/openapi-console.json b/sdk/typescript/openapi-console.json index c67de04c..354be8dd 100644 --- a/sdk/typescript/openapi-console.json +++ b/sdk/typescript/openapi-console.json @@ -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": { diff --git a/sdk/typescript/openapi-server.json b/sdk/typescript/openapi-server.json index 3ab5541b..8f4df48c 100644 --- a/sdk/typescript/openapi-server.json +++ b/sdk/typescript/openapi-server.json @@ -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": { diff --git a/tests/test_coordinator_client.py b/tests/test_coordinator_client.py index 1c49197c..b89b523e 100644 --- a/tests/test_coordinator_client.py +++ b/tests/test_coordinator_client.py @@ -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, ): diff --git a/tests/test_coordinator_quota_endpoint.py b/tests/test_coordinator_quota_endpoint.py deleted file mode 100644 index 24fd5422..00000000 --- a/tests/test_coordinator_quota_endpoint.py +++ /dev/null @@ -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) diff --git a/tests/test_coordinator_tools.py b/tests/test_coordinator_tools.py index a66f551c..9c4f9f3f 100644 --- a/tests/test_coordinator_tools.py +++ b/tests/test_coordinator_tools.py @@ -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 diff --git a/tests/test_spawn_quota.py b/tests/test_spawn_quota.py deleted file mode 100644 index c6f29052..00000000 --- a/tests/test_spawn_quota.py +++ /dev/null @@ -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 diff --git a/turnstone/api/console_schemas.py b/turnstone/api/console_schemas.py index 39114d49..be384c77 100644 --- a/turnstone/api/console_schemas.py +++ b/turnstone/api/console_schemas.py @@ -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. diff --git a/turnstone/api/console_spec.py b/turnstone/api/console_spec.py index 9fff5dbe..a5626e1b 100644 --- a/turnstone/api/console_spec.py +++ b/turnstone/api/console_spec.py @@ -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, diff --git a/turnstone/console/coordinator_client.py b/turnstone/console/coordinator_client.py index 0e294410..649d8169 100644 --- a/turnstone/console/coordinator_client.py +++ b/turnstone/console/coordinator_client.py @@ -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 diff --git a/turnstone/console/server.py b/turnstone/console/server.py index ad8e62c2..eff6828c 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -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, diff --git a/turnstone/core/session.py b/turnstone/core/session.py index ca2b577b..5933ddb7 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -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" ) diff --git a/turnstone/core/settings_registry.py b/turnstone/core/settings_registry.py index 91a23907..1ad53848 100644 --- a/turnstone/core/settings_registry.py +++ b/turnstone/core/settings_registry.py @@ -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} diff --git a/turnstone/core/spawn_quota.py b/turnstone/core/spawn_quota.py deleted file mode 100644 index 1a1c3875..00000000 --- a/turnstone/core/spawn_quota.py +++ /dev/null @@ -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, - ) diff --git a/turnstone/core/storage/migrations/versions/047_drop_coord_spawn_quota_settings.py b/turnstone/core/storage/migrations/versions/047_drop_coord_spawn_quota_settings.py new file mode 100644 index 00000000..496546f6 --- /dev/null +++ b/turnstone/core/storage/migrations/versions/047_drop_coord_spawn_quota_settings.py @@ -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