Look up accounting reservations by primary key from the usage trigger

CHANGELOG.md

@@ -20,6 +20,11 @@ SPDX-License-Identifier: Apache-2.0
 
 ## Unreleased
 
+- Starting a run no longer slows down as executions accumulate. The usage accounting that runs on every run start
+  and every change to quota-counted records read every reservation in the database, one per run ever started, and
+  now finds its reservation and allocation entries by primary key. Usage figures and limits are unchanged. New
+  migration `0031_operations_key_collation` replaces the two internal accounting functions in place and changes no
+  table, data or role; migrate the server explicitly before restarting it.
 - Studio has one place to set up AI: **Settings › AI models**. It shows provider cards, your AI connections, a
   readiness checklist with **Publish**, **Authorize** and **Test connection**, and the Weave AI card. The AI
   connection side sheet is the only AI connection form; it also opens from Build › Connections, an AI task and

docs/operations/upgrades.md

@@ -45,7 +45,7 @@ the desktop app never migrates a platform.
 | alpha13 | `0030_worker_presence` | No new schema migration. Workers retry explicit capacity rejections while they read task context or credentials, and an Agentic preparation timeout before provider execution counts as not started. Upgrade the Agentic and Files workers to 0.1.5 together with the server. |
 | alpha14 | `0030_worker_presence` | No new schema migration. Adds the detached Docker development platform (`weave platform up`); existing foreground installations keep working and are never converted. The Agentic and Files workers 0.1.6 pin this server version. |
 | alpha15 | `0030_worker_presence` | No new schema migration and no role change. Upgrade CLIs and SDKs together with the server, review native executor `capacity` and the legacy private-network settings before restarting, and finish runs that use text templates before any rollback to alpha14; see [what changes in alpha15](#what-changes-in-alpha15). The Agentic and Files workers 0.1.7 pin this server version. |
-| Unreleased | `0030_worker_presence` | No schema change. `weave platform up --username` grants the local owner roles (developer, deployer, operator, viewer, task_participant, tenant_admin, lumi_manager, execution_manager, worker_operator, deployment_reader, deployment_planner, deployment_approver, deployment_operator); existing accounts keep their roles until `weave platform user grant` adds more. `weave platform user` is now a command group; `weave platform user --username` works as before. |
+| Unreleased | `0031_operations_key_collation` | One new migration that only replaces two internal accounting functions: no table, data, or role changes. Explicitly migrate the server before restarting it. `weave platform up --username` grants the local owner roles (developer, deployer, operator, viewer, task_participant, tenant_admin, lumi_manager, execution_manager, worker_operator, deployment_reader, deployment_planner, deployment_approver, deployment_operator); existing accounts keep their roles until `weave platform user grant` adds more. `weave platform user` is now a command group; `weave platform user --username` works as before. |
 
 ![Schema, compatibility and execution acceptance gates](../diagrams/operations-upgrade.svg)
 
@@ -359,8 +359,8 @@ changes before you reopen traffic:
 
 ## What changes after alpha15
 
-These changes are not yet in a release. They add no schema migration (the head
-stays `0030_worker_presence`), no role, and no change to definition schemas or
+These changes are not yet in a release. They add one schema migration (the head
+becomes `0031_operations_key_collation`), no role, and no change to definition schemas or
 the stored digests of the built-in AI Connector, Action and task. Check them
 before you reopen traffic on a server that includes them:
 
@@ -411,6 +411,13 @@ before you reopen traffic on a server that includes them:
   A caller whose grant was revoked gets 403 (or 401) instead of the stored 201
   answer. Scripts that retry activations with a token whose grant was removed
   must get the grant back first.
+- **Starting a run no longer slows down as executions accumulate.** Migration
+  `0031_operations_key_collation` fixes the usage accounting that runs on
+  every run start and every change to quota-counted records: it now finds its
+  reservation and allocation entries by primary key instead of reading every
+  reservation in the database, which grew with each run ever started. Usage
+  figures and limits are unchanged. Apply the migration before you restart the
+  server; a server that includes it refuses to start on the previous schema.
 - **AI test times use the database clock.** A recorded AI connection test is
   dated with the database's clock, the same clock readiness uses to measure its
   age. Tests recorded by an earlier server keep their API-host timestamps.

migrations/versions/0031_operations_key_collation.py

@@ -0,0 +1,56 @@
+# Copyright 2026 Firefly Software Foundation.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Author: Firefly Software Foundation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Collation-safe accounting lookups so every charged write stays a primary-key lookup."""
+
+from alembic import op
+
+revision = "0031_operations_key_collation"
+down_revision = "0030_worker_presence"
+
+# The accounting trigger passes TG_TABLE_NAME (type name, collation "C") down the call chain, so
+# PL/pgSQL gives every text parameter and variable collation "C". Compared with a default-collation
+# key column, that rules out the primary key and turns each lookup into a scan of the whole table.
+PATCHES = (
+    (
+        "weave_reserve_refresh(text,uuid,boolean)",
+        "kind=kind_name AND resource_id=identifier",
+        'kind=(kind_name COLLATE "default") AND resource_id=identifier',
+        2,
+    ),
+    (
+        "weave_allocation_delta(text,jsonb,jsonb,text[])",
+        "WHERE relation=relation_name AND row_key=identity_key",
+        'WHERE relation=(relation_name COLLATE "default") AND row_key=(identity_key COLLATE "default")',
+        2,
+    ),
+)
+
+
+def upgrade():
+    for function, before, after, occurrences in PATCHES:
+        # CREATE OR REPLACE keeps the owner, grants, and SET options of the shared function.
+        op.execute(f"""DO $migration$ DECLARE body text; BEGIN
+          SELECT pg_get_functiondef('public.{function}'::regprocedure) INTO body;
+          IF (length(body)-length(replace(body,$p${before}$p$,'')))/length($p${before}$p$)<>{occurrences} THEN
+            RAISE EXCEPTION 'Unexpected accounting function {function}; verify predecessor'; END IF;
+          EXECUTE replace(body,$p${before}$p$,$p${after}$p$);
+        END $migration$""")
+    op.execute("UPDATE weave_schema_version SET version='0031_operations_key_collation'")
+
+
+def downgrade():
+    raise RuntimeError("Destructive downgrades are not supported")

src/firefly_weave/persistence/migrations.py

@@ -29,7 +29,7 @@
 
 from firefly_weave.settings import Settings, operations_from_env
 
-SCHEMA_VERSION = "0030_worker_presence"
+SCHEMA_VERSION = "0031_operations_key_collation"
 
 
 async def check_schema(connection: AsyncConnection) -> None:

tests/integration/operations_support.py

@@ -0,0 +1,123 @@
+# Copyright 2026 Firefly Software Foundation.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Author: Firefly Software Foundation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Independent logical-usage inventory shared by operations integration tests."""
+
+from collections import defaultdict
+from uuid import UUID
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncConnection
+
+ZERO = UUID(int=0)
+
+
+async def recompute_usage(connection: AsyncConnection, tenant: UUID, project: UUID) -> dict[tuple[UUID, str], int]:
+    """Inventory charged tables and retained control provenance independently of counters."""
+    tables = (
+        (
+            await connection.execute(
+                text(
+                    "SELECT c.relname FROM pg_trigger t JOIN pg_class c ON c.oid=t.tgrelid "
+                    "JOIN pg_namespace n ON n.oid=c.relnamespace WHERE t.tgname='operation_usage' "
+                    "AND n.nspname='public' "
+                    "ORDER BY c.relname"
+                )
+            )
+        )
+        .scalars()
+        .all()
+    )
+    totals = defaultdict(int)
+    params = {"tenant": tenant, "project": project}
+    for table in tables:
+        projection = f"public.weave_{table}_account(r)" if table in ("runs", "run_events") else "to_jsonb(r)"
+        rows = (
+            await connection.execute(
+                text(
+                    f"SELECT item,public.weave_row_metrics('{table}',item) AS metrics FROM "
+                    f"(SELECT {projection} AS item FROM {table} r) q WHERE "
+                    "coalesce(item->>'tenant_id',"
+                    "item->'event'->'scope'->>'tenant_id')=cast(cast(:tenant AS uuid) AS text) AND "
+                    "coalesce(item->>'project_id',"
+                    "item->'event'->'scope'->>'project_id')=cast(cast(:project AS uuid) AS text)"
+                ),
+                params,
+            )
+        ).mappings()
+        for row in rows:
+            for metric, amount in row["metrics"].items():
+                dimension = (
+                    ZERO
+                    if metric in ("ordinary_bytes", "runs_retained", "debug_rows")
+                    else UUID(row["item"].get("environment_id", str(ZERO)))
+                )
+                if metric == "outbox_pending" and await connection.scalar(
+                    text(
+                        "SELECT EXISTS(SELECT 1 FROM operation_control_deliveries WHERE tenant_id=:tenant "
+                        "AND project_id=:project AND id=:id)"
+                    ),
+                    {**params, "id": row["item"]["id"]},
+                ):
+                    metric, dimension = "control_outbox_pending", ZERO
+                totals[dimension, metric] += int(amount)
+    allocation = (
+        (
+            await connection.execute(
+                text(
+                    "SELECT coalesce(sum(control_bytes),0) AS spent,count(*) AS rows "
+                    "FROM operation_allocations WHERE tenant_id=:tenant AND project_id=:project"
+                ),
+                params,
+            )
+        )
+        .mappings()
+        .one()
+    )
+    totals[ZERO, "ordinary_bytes"] -= int(allocation["spent"])
+    totals[ZERO, "control_bytes"] += int(allocation["spent"]) + 512 * allocation["rows"]
+    reservation = (
+        (
+            await connection.execute(
+                text(
+                    "SELECT coalesce(sum(amount),0) AS amount,"
+                    "coalesce(sum(outbox_slots),0) AS slots,count(*) AS rows FROM operation_reservations "
+                    "WHERE tenant_id=:tenant AND project_id=:project"
+                ),
+                params,
+            )
+        )
+        .mappings()
+        .one()
+    )
+    totals[ZERO, "ordinary_bytes"] += 512 * reservation["rows"]
+    totals[ZERO, "control_reserved"] += int(reservation["amount"])
+    totals[ZERO, "control_outbox_reserved"] += int(reservation["slots"])
+    return {key: value for key, value in totals.items() if value}
+
+
+async def assert_usage(connection: AsyncConnection, ids: dict[str, UUID]) -> None:
+    expected = await recompute_usage(connection, ids["tenant"], ids["project"])
+    actual = (
+        await connection.execute(
+            text(
+                "SELECT environment_id,metric,value FROM operation_usage "
+                "WHERE tenant_id=:tenant AND project_id=:project AND value<>0"
+            ),
+            ids,
+        )
+    ).all()
+    assert {(environment, metric): value for environment, metric, value in actual} == expected

tests/integration/test_reservation_lookup.py

@@ -0,0 +1,136 @@
+# Copyright 2026 Firefly Software Foundation.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Author: Firefly Software Foundation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Trigger-side reservation lookups stay primary-key lookups however many reservations exist."""
+
+import json
+from uuid import UUID, uuid4
+
+import pytest
+import test_waits as wait_tests
+from operations_support import assert_usage
+from sqlalchemy import text
+
+pytestmark = pytest.mark.integration
+author = wait_tests.author
+waiting_run = wait_tests.waiting_run
+
+# Tables whose primary key the accounting trigger reaches with a value derived from TG_TABLE_NAME.
+CONTROL = "operation_control_context"
+LOOKUPS = {
+    "operation_reservations": ("operation_reservations_pkey", ("kind", "resource_id")),
+    "operation_allocations": ("operation_allocations_pkey", ("relation", "row_key")),
+}
+
+
+def scans(plan, found):
+    if plan.get("Relation Name") in (*LOOKUPS, CONTROL) and "Scan" in plan["Node Type"]:
+        found.append(plan)
+    for child in plan.get("Plans", ()):
+        scans(child, found)
+    return found
+
+
+async def explained_trigger(owner, statement, params):
+    """Run one statement and return every plan auto_explain reports for the trigger's nested queries."""
+    notices = []
+    raw = (await (await owner.connection()).get_raw_connection()).driver_connection
+    listener = lambda _, message: notices.append(message.message)  # noqa: E731
+    raw.add_log_listener(listener)
+    try:
+        for setting in (
+            "LOAD 'auto_explain'",
+            "SET LOCAL auto_explain.log_min_duration=0",
+            "SET LOCAL auto_explain.log_nested_statements=on",
+            "SET LOCAL auto_explain.log_format=json",
+            "SET LOCAL auto_explain.log_level=notice",
+            "SET LOCAL client_min_messages=notice",
+        ):
+            await owner.execute(text(setting))
+        await owner.execute(text(statement), params)
+        await owner.execute(text("SET LOCAL auto_explain.log_min_duration=-1"))
+    finally:
+        raw.remove_log_listener(listener)
+    return [json.loads(notice[notice.index("{") :]) for notice in notices if "plan:" in notice]
+
+
+async def test_run_insert_uses_reservation_primary_key_among_thousands(waiting_run, access_db, author):
+    ids = {"tenant": author[2].tenant_id, "project": author[2].project_id}
+    async with access_db[1].begin() as owner:
+        # Another project's reservations: the lookup must not cost time proportional to them.
+        other_tenant, other_project = uuid4(), uuid4()
+        await owner.execute(text("INSERT INTO tenants VALUES(:id,'filler')"), {"id": other_tenant})
+        await owner.execute(
+            text("INSERT INTO projects VALUES(:id,:tenant,'filler')"), {"id": other_project, "tenant": other_tenant}
+        )
+        await owner.execute(
+            text(
+                "INSERT INTO operation_reservations SELECT :tenant,:project,"
+                "(ARRAY['run','provider','binding','webhook'])[1+n%4],gen_random_uuid(),0,0 "
+                "FROM generate_series(1,6000) n"
+            ),
+            {"tenant": other_tenant, "project": other_project},
+        )
+        await owner.execute(
+            text(
+                "INSERT INTO operation_allocations SELECT :tenant,:project,"
+                "(ARRAY['runs','run_events','task_intents'])[1+n%3],md5(n::text),1 "
+                "FROM generate_series(1,6000) n"
+            ),
+            {"tenant": other_tenant, "project": other_project},
+        )
+        await owner.execute(text("ANALYZE operation_reservations"))
+        await owner.execute(text("ANALYZE operation_allocations"))
+        await assert_usage(owner, ids)
+        original = (
+            await owner.execute(
+                text("SELECT (amount,outbox_slots) FROM operation_reservations WHERE kind='run' AND resource_id=:id"),
+                {"id": UUID(waiting_run)},
+            )
+        ).scalar_one()
+    copy = uuid4()
+    async with access_db[1].begin() as owner:
+        plans = await explained_trigger(
+            owner,
+            "INSERT INTO runs SELECT (jsonb_populate_record(r,jsonb_build_object('id',cast(:copy AS uuid)))).* "
+            "FROM runs r WHERE id=cast(:run AS uuid)",
+            {"copy": copy, "run": waiting_run},
+        )
+        found = [scan for plan in plans for scan in scans(plan["Plan"], [])]
+        assert {*LOOKUPS, CONTROL} <= {scan["Relation Name"] for scan in found}, plans
+        # The control context is empty outside a reserved control, so only its key collation is meaningful.
+        assert not [scan for scan in found if 'COLLATE "C"' in scan.get("Index Cond", "") + scan.get("Filter", "")]
+        for scan in (scan for scan in found if scan["Relation Name"] in LOOKUPS):
+            index, columns = LOOKUPS[scan["Relation Name"]]
+            condition = scan.get("Index Cond", "")
+            assert scan["Node Type"] in {"Index Scan", "Index Only Scan"}, scan
+            assert scan["Index Name"] == index, scan
+            assert all(f"{column} =" in condition for column in columns), scan
+        reserved = dict(
+            (
+                await owner.execute(
+                    text(
+                        "SELECT resource_id,(amount,outbox_slots) FROM operation_reservations "
+                        "WHERE kind='run' AND resource_id IN (:copy,:run)"
+                    ),
+                    {"copy": copy, "run": UUID(waiting_run)},
+                )
+            ).all()
+        )
+        assert reserved[UUID(waiting_run)] == original
+        amount, slots = reserved[copy]
+        assert amount > 16777216 and slots == 64
+        await assert_usage(owner, ids)