From dc9811f5bd45aa2b984e0364412ea06c83f68a45 Mon Sep 17 00:00:00 2001 From: Toby Martin Date: Fri, 21 Aug 2026 10:48:37 +1000 Subject: [PATCH] feat(ai): add audited budget reservation reclamation --- Cargo.lock | 4 +- Cargo.toml | 2 +- .../graphql-orm-ai-tool-profiles/CHANGELOG.md | 17 + .../graphql-orm-ai-tool-profiles/Cargo.toml | 2 +- .../graphql-orm-ai-tool-profiles/MIGRATION.md | 17 + crates/graphql-orm-ai-tool-profiles/README.md | 2 +- .../graphql-orm-ai-tool-profiles/src/error.rs | 9 + crates/graphql-orm-ai/CHANGELOG.md | 60 + crates/graphql-orm-ai/Cargo.toml | 2 +- crates/graphql-orm-ai/MIGRATION.md | 102 ++ crates/graphql-orm-ai/README.md | 18 +- .../docs/implementation-status.md | 6 +- .../graphql-orm-ai/docs/usage-and-budgets.md | 81 +- crates/graphql-orm-ai/src/configuration.rs | 237 +++ crates/graphql-orm-ai/src/orm_budget.rs | 169 ++- .../graphql-orm-ai/src/orm_configuration.rs | 1338 ++++++++++++++++- crates/graphql-orm-ai/src/orm_coordinator.rs | 210 ++- .../src/orm_supervised_coordinator.rs | 92 +- crates/graphql-orm-ai/src/persistence.rs | 13 +- crates/graphql-orm-ai/src/provider_calls.rs | 289 +++- crates/graphql-orm-ai/src/run_state.rs | 1 + crates/graphql-orm-ai/tests/graphql_naming.rs | 8 + crates/graphql-orm-ai/tests/schema_module.rs | 7 +- docs/reference/workspace-packages.md | 4 +- 24 files changed, 2633 insertions(+), 57 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5771df23..57d1b43b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3104,7 +3104,7 @@ dependencies = [ [[package]] name = "graphql-orm-ai" -version = "0.83.0" +version = "0.84.0" dependencies = [ "agql-auth", "async-graphql", @@ -3136,7 +3136,7 @@ dependencies = [ [[package]] name = "graphql-orm-ai-tool-profiles" -version = "0.6.0" +version = "0.7.0" dependencies = [ "async-graphql", "async-graphql-parser", diff --git a/Cargo.toml b/Cargo.toml index 40cb50ad..2b1f1ab5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,7 +38,7 @@ futures = "0.3" getrandom = "0.3" graphql-composition = "=0.12.2" graphql-orm = { path = "crates/graphql-orm", version = "0.23.0", default-features = false } -graphql-orm-ai-tool-profiles = { path = "crates/graphql-orm-ai-tool-profiles", version = "0.6.0" } +graphql-orm-ai-tool-profiles = { path = "crates/graphql-orm-ai-tool-profiles", version = "0.7.0" } graphql-orm-backup = { path = "crates/graphql-orm-backup", version = "0.7.1", default-features = false } graphql-orm-operation-catalog = { path = "crates/graphql-orm-operation-catalog", version = "0.3.0" } graphql-orm-router-protocol = { path = "crates/graphql-orm-router-protocol", version = "0.2.1" } diff --git a/crates/graphql-orm-ai-tool-profiles/CHANGELOG.md b/crates/graphql-orm-ai-tool-profiles/CHANGELOG.md index 141898c1..88ea7501 100644 --- a/crates/graphql-orm-ai-tool-profiles/CHANGELOG.md +++ b/crates/graphql-orm-ai-tool-profiles/CHANGELOG.md @@ -10,6 +10,23 @@ supersedes: [] # Changelog +## [0.7.0] - 2026-08-21 + +### Added + +- `AiError::PreTransportBudgetDenied` is the closed execution-boundary signal + for an atomic budget refusal proven to occur before provider dispatch and + after any created reservation was released. It retains the existing public + `AI_BUDGET_DENIED` code while preventing generic tool-loop budget limits from + being mistaken for proof that provider transport never occurred. + +### Breaking + +- `AiError` gained `PreTransportBudgetDenied`. Although the enum is + non-exhaustive, in-crate and deliberately exhaustive consumers must handle + the new variant. `BudgetDenied` remains the generic limit error and carries + no transport-absence proof. + ## [0.6.0] - 2026-08-16 ### Added diff --git a/crates/graphql-orm-ai-tool-profiles/Cargo.toml b/crates/graphql-orm-ai-tool-profiles/Cargo.toml index 39440847..5c8da7b4 100644 --- a/crates/graphql-orm-ai-tool-profiles/Cargo.toml +++ b/crates/graphql-orm-ai-tool-profiles/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "graphql-orm-ai-tool-profiles" -version = "0.6.0" +version = "0.7.0" edition = "2024" authors = ["Toby Martin "] description = "Backend-neutral GraphQL AI tool profile compiler and manifest contracts" diff --git a/crates/graphql-orm-ai-tool-profiles/MIGRATION.md b/crates/graphql-orm-ai-tool-profiles/MIGRATION.md index 2797ef8a..80881eac 100644 --- a/crates/graphql-orm-ai-tool-profiles/MIGRATION.md +++ b/crates/graphql-orm-ai-tool-profiles/MIGRATION.md @@ -10,6 +10,23 @@ supersedes: [] # Migration Guide +## 0.6.0 to 0.7.0: proof-bearing pre-transport budget denial + +Adopt `graphql-orm-ai-tool-profiles` 0.7.0 with `graphql-orm-ai` 0.84.0 from +one reviewed full monorepo revision. + +`AiError` gains `PreTransportBudgetDenied`. Only the provider-execution +boundary may return it, and only after proving provider dispatch was never +attempted and any created reservation was durably released. A budget or rule +limit reached during a provider/tool loop remains `BudgetDenied`; it must not +be reclassified as a certain local refusal. Both variants intentionally expose +the same stable public error code, `AI_BUDGET_DENIED`. + +The enum was already non-exhaustive. Update any deliberately exhaustive +internal matches. There is no schema, database, data, GraphQL SDL, manifest, +capability, fingerprint, protected-content, credential, or AI schema-module +migration from this package change. + ## 0.5.0 to 0.6.0: compact discovery and query-plan wire v3 Adopt `graphql-orm-ai-tool-profiles` 0.6.0 and `graphql-orm-ai` 0.81.0 from one diff --git a/crates/graphql-orm-ai-tool-profiles/README.md b/crates/graphql-orm-ai-tool-profiles/README.md index 403f7891..5353c35a 100644 --- a/crates/graphql-orm-ai-tool-profiles/README.md +++ b/crates/graphql-orm-ai-tool-profiles/README.md @@ -24,7 +24,7 @@ are separate runtime decisions and must remain default-deny. ```toml [dependencies] -graphql-orm-ai-tool-profiles = { git = "https://github.com/Dastari/graphql-orm.git", rev = "", version = "0.6.0" } +graphql-orm-ai-tool-profiles = { git = "https://github.com/Dastari/graphql-orm.git", rev = "", version = "0.7.0" } serde_json = "1" ``` diff --git a/crates/graphql-orm-ai-tool-profiles/src/error.rs b/crates/graphql-orm-ai-tool-profiles/src/error.rs index ef6e4ced..935bbc33 100644 --- a/crates/graphql-orm-ai-tool-profiles/src/error.rs +++ b/crates/graphql-orm-ai-tool-profiles/src/error.rs @@ -31,6 +31,14 @@ pub enum AiError { /// No applicable atomic budget had enough capacity for the operation. #[error("AI budget denied")] BudgetDenied, + /// Provider execution was denied by an atomic budget before dispatch. + /// + /// This is an execution-boundary proof, not a generic budget error. It may + /// be returned only when provider transport was never attempted and no + /// budget reservation remains held. Budget limits reached during a + /// provider/tool loop must use [`Self::BudgetDenied`] instead. + #[error("AI budget denied")] + PreTransportBudgetDenied, /// Input failed a public schema contract. #[error("invalid AI input: {0}")] InvalidInput(String), @@ -66,6 +74,7 @@ impl AiError { Self::RecentMfaRequired => "AI_RECENT_MFA_REQUIRED", Self::EgressDenied => "AI_EGRESS_DENIED", Self::BudgetDenied => "AI_BUDGET_DENIED", + Self::PreTransportBudgetDenied => "AI_BUDGET_DENIED", Self::InvalidInput(_) => "AI_INVALID_INPUT", Self::ReauthorizationFailed => "AI_REAUTHORIZATION_FAILED", Self::ToolExecutionFailed => "AI_TOOL_EXECUTION_FAILED", diff --git a/crates/graphql-orm-ai/CHANGELOG.md b/crates/graphql-orm-ai/CHANGELOG.md index 972b9276..db2b2bfc 100644 --- a/crates/graphql-orm-ai/CHANGELOG.md +++ b/crates/graphql-orm-ai/CHANGELOG.md @@ -18,6 +18,66 @@ checkpoint facts. For the current workspace baseline and active gates, use the [implementation status](docs/implementation-status.md) and the central [AI production-readiness plan](../../docs/plans/active/ai-production-readiness/README.md). +## [Unreleased] + +Persistent schema module: **0.63.0**. + +### Added + +- `aiBudgetScopeCapacity` reports, for one exact scope, each budget policy's + current-period reserved and committed amounts beside its ceilings, counts of + unresolved reservations, and a bounded oldest-first list of them with expiry, + owning-run terminality, CAS version, and reclaimability. Reserved capacity + counts against a ceiling exactly like committed usage, so a host can now + alarm on stranded reservations instead of discovering them as a total refusal + to serve. Authorized by the existing `ReadBudgetPolicies` action. +- `reclaimAiBudgetReservation` resolves one stranded budget reservation whose + capacity could previously never leave the reserved column. It requires the + new `AiConfigurationAction::ManageBudgetReclamation`, recent MFA, an exact + CAS version, an expiry already past the deployment's `minimum_expired_age`, + and an owning run in a durable terminal state holding no lease. It commits + the reservation's own reserved amounts as authoritative usage and appends one + usage fact and one redacted audit fact in the same transaction. It never + releases capacity: an `uncertain` or `reserved` reservation carries no proof + that the provider was not reached, so the only safe resolution is the + conservative one, which can over-count and never under-counts. +- `AiBudgetReclamationLimits` and + `OrmAiConfigurationService::with_budget_reservation_reclamation` are the + deployment opt-in for the reclamation surface. Without them, capacity + reporting still works and every reservation reports `reclaimable: false`. + +### Changed + +- **A proven pre-transport budget denial is no longer reported as provider + uncertainty.** `AiError::PreTransportBudgetDenied` is produced only when the + atomic reservation was refused before dispatch, or an already-created + reservation was durably released before dispatch. Such a run now terminates + `Failed` with outcome code `provider_budget_denied` instead of + `RecoveryRequired` with `provider_turn_uncertain`, and + `provider_budget_denied` joins the retryable failure allowlist so the + terminal-event failure record reports `AiRunRetryAdmission::Allowed`. The + supervised coordinator makes the same distinction. Generic + `AiError::BudgetDenied`, including a limit reached inside a dynamic tool + loop, remains uncertain and can never claim transport absence. +- A provider call whose post-reservation authorization binding fails now + releases its reservation instead of leaving capacity held until the policy + period rolled. Nothing had been dispatched, so the release is provable. + +### Breaking + +- `AiConfigurationAction` gained `ManageBudgetReclamation`. Exhaustive matches + in host access policies must handle it; a host that does not recognize it + must deny. +- `AiError` gained the proof-bearing `PreTransportBudgetDenied` variant through + `graphql-orm-ai-tool-profiles` 0.7.0. It shares the existing public code with + `BudgetDenied`, but exhaustive internal matches must preserve their distinct + transport semantics. +- The persistent schema module advances to 0.63.0. It adds no entity, column, + or constraint. It makes `scope_kind`, `scope_id`, `tenant_id`, and + `expires_at` available to typed internal predicates and adds one composite + scope/tenant/state/expiry index, so stranded-reservation reporting is a + bounded indexed read. + ## [0.83.0] - 2026-08-21 Persistent schema module: **0.62.0**. diff --git a/crates/graphql-orm-ai/Cargo.toml b/crates/graphql-orm-ai/Cargo.toml index 8df49548..ea4d0ac4 100644 --- a/crates/graphql-orm-ai/Cargo.toml +++ b/crates/graphql-orm-ai/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "graphql-orm-ai" -version = "0.83.0" +version = "0.84.0" edition = "2024" authors = ["Toby Martin "] description = "Project-agnostic AI agent runtime for graphql-orm applications" diff --git a/crates/graphql-orm-ai/MIGRATION.md b/crates/graphql-orm-ai/MIGRATION.md index 812d2c91..0d1397e5 100644 --- a/crates/graphql-orm-ai/MIGRATION.md +++ b/crates/graphql-orm-ai/MIGRATION.md @@ -19,6 +19,108 @@ they describe. For the current workspace baseline and active delivery gates, use [implementation status](docs/implementation-status.md) and the central [AI production-readiness plan](../../docs/plans/active/ai-production-readiness/README.md). +## Unreleased: budget reclamation and pre-transport denial (crate 0.83.0 to 0.84.0; schema 0.62.0 to 0.63.0) + +### Schema module + +The AI schema module advances **0.62.0 to 0.63.0**. It adds **no entity, no +column, and no constraint**, and needs **no data migration or backfill**. +`graphql_orm_ai_budget_reservations` exposes `scope_kind`, `scope_id`, +`tenant_id`, and `expires_at` to typed internal predicates and gains the composite index +`idx_graphql_orm_ai_budget_reservations_scope_state` +(`scope_kind, scope_id, tenant_id, state, expires_at`), so the new stranded-reservation +report is a bounded indexed read rather than a table scan. Existing rows and +events remain readable at the previous module version and after the upgrade. +Apply and verify the module before serving traffic. + +### Source-breaking changes + +`AiConfigurationAction` gained the variant `ManageBudgetReclamation`. Any +exhaustive `match` in a host `AiConfigurationAccessPolicy` must handle it. +Authorize it only for the administrators you would trust to charge an +unprovable provider turn to a budget; a host that does not want the surface +returns `false` and it stays closed. + +`AiConfigurationService` gained `budget_scope_capacity` and +`reclaim_budget_reservation`. Both have fail-closed default implementations +that return `AiError::InvalidConfiguration`, so an existing custom +implementation still compiles and does not silently gain the surface. + +### New public API + +- `AiBudgetScopeCapacityView`, `AiBudgetPolicyCapacityView`, and + `AiBudgetReservationCapacityView` are the redacted capacity views. They carry + capacity accounting, reservation state, expiry, owning-run linkage and CAS + versions only; never a prompt, transcript, provider payload, principal + identity, or credential. `reclaimable` identifies a deployment/time/run + candidate only; mutation authorization, recent MFA, CAS, scope and stored + graph integrity are rechecked separately. +- `ReclaimAiBudgetReservationInput { scope, reservation_id, expected_version }`. +- `AiBudgetReclamationLimits::new(minimum_expired_age, maximum_reservation_scan)` + and `OrmAiConfigurationService::with_budget_reservation_reclamation`. + +### New GraphQL + +`AiConfigurationQueryRoot` gains `aiBudgetScopeCapacity(scope)`, authorized by +the existing `ReadBudgetPolicies` action. `AiConfigurationMutationRoot` gains +`reclaimAiBudgetReservation(input)`, authorized by `ManageBudgetReclamation` +plus recent MFA plus the deployment opt-in. No existing field changed. + +To enable reclamation: + +```rust,ignore +let configuration = OrmAiConfigurationService::new(/* ... */) + .with_budget_policy_management(policy_limits) + .with_budget_reservation_reclamation(AiBudgetReclamationLimits::new( + time::Duration::hours(6), + 200, + )?); +``` + +Without that call, `aiBudgetScopeCapacity` still works and every reservation +reports `reclaimable: false`, while `reclaimAiBudgetReservation` fails closed +as invalid configuration. + +### Behavioural changes with no API change + +A run returning the proof-bearing `AiError::PreTransportBudgetDenied` now +terminates `Failed` with outcome code `provider_budget_denied`. The executor +may produce that variant only when reservation failed before dispatch or when +an already-created reservation was durably released before dispatch. It +previously terminated `RecoveryRequired` with `provider_turn_uncertain`, which +told users a proven local refusal could not be confirmed and made the run +permanently unretryable. `provider_budget_denied` is on the retryable failure +allowlist, so `AiRunFailure.admission` is `AiRunRetryAdmission::Allowed` and a +client may author a new run for the same durable user message once capacity +exists. A client that keys UI text off `provider_turn_uncertain` for this case +must move it to the new code. The supervised coordinator makes the same +distinction. Generic `AiError::BudgetDenied`, including post-transport dynamic +tool-call and rule ceilings, remains uncertain and must not use this path. + +A provider call whose post-reservation authorization binding fails now releases +its reservation. Nothing had been dispatched, so the release is provable. + +### What deliberately did not change + +Reclamation commits; it never releases. An `uncertain` or `reserved` +reservation carries no durable proof that the provider was not reached, so +releasing it would fabricate an absence proof. Committing the held estimate can +only over-count. + +Reclamation therefore **does not create headroom**: the reserved column falls +by exactly the amount the committed column rises, and `reserved + committed` +against the ceiling is unchanged. A deployment whose ceiling is already +exhausted by stranded reservations raises or replaces the policy through +`upsertAiBudgetPolicy`. The value of reclamation is that held capacity becomes +accountable, reportable, and finite instead of permanently unreachable, and +that the condition is now observable before it becomes an outage. + +Reclamation is not automatic. Expired-lease recovery and every other +maintenance pass are unchanged. Automating a commit would free no headroom +while adding an unattended writer of authoritative usage facts attributed to an +absent principal; the decision to charge an unprovable turn stays with an +authorized, MFA-current, audited human. + ## 0.82.0 to 0.83.0: settled retained Codex interruption Adopt `graphql-orm-ai` 0.83.0 at one reviewed full monorepo revision. diff --git a/crates/graphql-orm-ai/README.md b/crates/graphql-orm-ai/README.md index b58a0d3b..63f4ed7f 100644 --- a/crates/graphql-orm-ai/README.md +++ b/crates/graphql-orm-ai/README.md @@ -28,7 +28,7 @@ for AI, ORM, storage, backup, and tool-profile packages: ```toml [dependencies] -graphql-orm-ai = { git = "https://github.com/Dastari/graphql-orm.git", rev = "", version = "0.82.0", default-features = false, features = ["sqlite"] } +graphql-orm-ai = { git = "https://github.com/Dastari/graphql-orm.git", rev = "", version = "0.84.0", default-features = false, features = ["sqlite"] } ``` Exactly one persistence backend is required: `sqlite` (default), `postgres`, @@ -173,6 +173,22 @@ Use `scripts/check-ai-provider-lanes.sh` from the repository root to verify one feature at a time; provider feature unification is not required for a valid adapter build. +## Budget capacity and stranded reservations + +Reserved capacity counts against a budget ceiling exactly like committed +usage, so a reservation that never reconciles consumes the ceiling for the rest +of its policy period. `aiBudgetScopeCapacity` reports per-policy reserved and +committed amounts, ceilings, and a bounded list of unresolved reservations +under `ReadBudgetPolicies`. `reclaimAiBudgetReservation` resolves one expired +reservation whose owning run is terminal, under `ManageBudgetReclamation`, +recent MFA, an exact CAS version, and the +`with_budget_reservation_reclamation` deployment opt-in. It commits the held +estimate as authoritative usage rather than releasing it, because an +unreconciled reservation carries no proof that the provider was not reached. +A denial at reservation is pre-transport and certain: the run fails with +`provider_budget_denied` and stays retryable. See the +[usage and budgets guide](docs/usage-and-budgets.md). + ## Reasoning effort profiles `ModelReasoningEffort` is the closed provider-neutral selection: diff --git a/crates/graphql-orm-ai/docs/implementation-status.md b/crates/graphql-orm-ai/docs/implementation-status.md index b595c3a4..dc3b844e 100644 --- a/crates/graphql-orm-ai/docs/implementation-status.md +++ b/crates/graphql-orm-ai/docs/implementation-status.md @@ -10,9 +10,9 @@ supersedes: [] # Implementation Status -`graphql-orm-ai` is at crate version `0.81.0` with AI schema module -`0.60.0`. It uses workspace `graphql-orm` `0.23.0`, backend-neutral -`graphql-orm-ai-tool-profiles` `0.6.0`, and external `agql-auth` +`graphql-orm-ai` is at crate version `0.84.0` with AI schema module +`0.63.0`. It uses workspace `graphql-orm` `0.23.0`, backend-neutral +`graphql-orm-ai-tool-profiles` `0.7.0`, and external `agql-auth` `0.15.0` at `e841ffd382082ad7419be259fe957f949b956ff7`. The active work order, dependencies, and exit gates are maintained in the diff --git a/crates/graphql-orm-ai/docs/usage-and-budgets.md b/crates/graphql-orm-ai/docs/usage-and-budgets.md index cbc58d6c..c38652ac 100644 --- a/crates/graphql-orm-ai/docs/usage-and-budgets.md +++ b/crates/graphql-orm-ai/docs/usage-and-budgets.md @@ -34,8 +34,71 @@ the same transaction. The usage fact has a unique reservation ID, so an exact idempotent replay returns the prior result without duplicating usage. `ReleaseUnused` is permitted only when transport provably did not occur and -appends no usage. `MarkUncertain` retains capacity and appends no usage; -privileged recovery remains a separately gated future surface. +appends no usage. `MarkUncertain` retains capacity and appends no usage; the +privileged recovery surface is described under +[Stranded reservations](#stranded-reservations-and-privileged-reclamation). + +A denial at reservation happens strictly before the transport boundary. It is +therefore a certain, local refusal, not provider uncertainty: it consumes no +provider turn and leaves no reservation held. The coordinator closes such a run +as terminal `Failed` with outcome code `provider_budget_denied`, and the +run-failure record admits a retry once capacity exists. It never closes the run +as `RecoveryRequired`/`provider_turn_uncertain`, which would tell a user that a +proven refusal could not be confirmed and would permanently refuse retry. + +## Stranded reservations and privileged reclamation + +Reserved capacity counts against a policy ceiling exactly like committed usage. +A reservation that never reconciles therefore consumes the ceiling for the rest +of its policy period, and if enough of them accumulate the deployment starts +refusing every new provider call. + +Two reservation states can strand: + +- `uncertain`, when the worker died after the transport boundary; and +- `reserved`, when the worker died between the reservation transaction and the + transport boundary. + +Neither carries a durable proof that the provider was not reached, so neither +may be released. `expires_at` bounds how long a reservation could still belong +to a live provider call; it does not prove anything about transport. + +`aiBudgetScopeCapacity` reports, for one exact scope under +`ReadBudgetPolicies`, each policy's current-period reserved and committed +amounts beside its ceilings, counts of unresolved reservations, and a bounded +oldest-first list carrying each reservation's state, expiry, owning-run +terminality, CAS version, and whether it meets the deployment and durable +time/run conditions to be a reclamation candidate. The mutation still rechecks +current authorization, recent MFA, exact CAS, scope, and stored-graph +integrity. Every count is a lower bound when `truncated` is set. Alarm on a +rising unresolved count rather than on the eventual refusal to serve. + +`reclaimAiBudgetReservation` resolves one exact reservation. It requires +`ManageBudgetReclamation` for the exact scope, a current user principal with +recent MFA, the deployment-owned +`OrmAiConfigurationService::with_budget_reservation_reclamation` opt-in, an +exact CAS version, an expiry that has already passed by the deployment's +`minimum_expired_age`, and an owning run that reached a durable terminal state +holding no lease. It then commits the reservation's own reserved amounts as +authoritative usage, appends one usage fact and one redacted audit fact, and +CAS-updates the reservation to `committed`, all in one state-machine +transaction. + +That resolution is conservative by construction: it charges the estimate that +was already being held, so it can only over-count. It also does not create +headroom. The reserved column falls by exactly the amount the committed column +rises, and `reserved + committed` against the ceiling is unchanged. What it +does is make held capacity accountable, reportable, and finite instead of +permanently unreachable. A deployment that needs headroom back raises or +replaces the policy through `upsertAiBudgetPolicy`; the crate will not +manufacture an absence proof to release capacity it cannot account for. + +Reclamation is **not** performed automatically by expired-lease recovery or any +other maintenance pass. Automation would buy nothing operationally, because +committing frees no headroom, while it would add an unattended writer of +authoritative usage facts attributed to a principal who is not present. The +decision to charge an unprovable turn stays with an authorized, MFA-current, +audited human. `input_tokens` is the provider-reported total. `cached_input_tokens` is a validated subset, not an additional amount. Deployment pricing can calculate @@ -46,8 +109,10 @@ authoritative total. Budget policies are managed only through `AiConfigurationQueryRoot` and `AiConfigurationMutationRoot`; private generated CRUD is not an application -surface. The host authorizes `ReadBudgetPolicies` and `ManageBudgetPolicies` -separately. Mutations also require a current user principal with recent MFA. +surface. The host authorizes `ReadBudgetPolicies`, `ManageBudgetPolicies`, and +`ManageBudgetReclamation` separately; a host that does not recognize an action +must return `false`. Mutations also require a current user principal with +recent MFA. Before enabling mutations, the deployment must call `OrmAiConfigurationService::with_budget_policy_management` with @@ -186,9 +251,11 @@ failure closes the query. ## Migration, backup, and restore The usage ledger was introduced by schema module `0.17.0`; current deployments -apply module `0.54.0`, whose append-only immutable pricing catalog includes -defaulted web/file-search per-call rates. Keep workers, configuration writes, -and readers closed during each managed migration. Unsupported legacy private +apply module `0.62.0`, which adds a scope/tenant/state/expiry index to +`graphql_orm_ai_budget_reservations` so stranded-reservation reporting is a +bounded indexed read. Module `0.54.0`'s append-only immutable pricing catalog +includes defaulted web/file-search per-call rates. Keep workers, configuration +writes, and readers closed during each managed migration. Unsupported legacy private usage rows must be proven from complete committed reservation evidence or removed; never fabricate a binding. Existing committed reservations are not automatically treated as historical diff --git a/crates/graphql-orm-ai/src/configuration.rs b/crates/graphql-orm-ai/src/configuration.rs index 80f25cb5..12fe8863 100644 --- a/crates/graphql-orm-ai/src/configuration.rs +++ b/crates/graphql-orm-ai/src/configuration.rs @@ -34,6 +34,15 @@ pub enum AiConfigurationAction { ReadBudgetPolicies, /// Create, alter, enable, or disable budget policies. ManageBudgetPolicies, + /// Resolve a stranded budget reservation whose capacity can otherwise + /// never leave the reserved column. + /// + /// This is deliberately separate from [`Self::ManageBudgetPolicies`]. It + /// does not change a ceiling; it converts held capacity into an + /// authoritative usage fact for one exact reservation. A host that does + /// not recognize this action must return `false` and keep the surface + /// closed. + ManageBudgetReclamation, /// Read immutable provider/model pricing versions. ReadPricingCatalog, /// Append immutable provider/model pricing versions. @@ -264,6 +273,151 @@ pub struct AiBudgetPolicyView { pub updated_at: i64, } +/// Observable capacity of one budget policy for its current period. +/// +/// Reserved and committed amounts are both counted against the policy ceiling. +/// Reserved capacity that never reconciles is therefore indistinguishable from +/// spend, which is why the stranded-reservation counts below sit beside it. +#[derive(Clone, Debug, SimpleObject)] +#[cfg_attr(feature = "graphql-case-pascal", graphql(rename_fields = "PascalCase"))] +pub struct AiBudgetPolicyCapacityView { + /// Policy ID. + pub policy_id: Uuid, + /// Stable reset interval. + pub interval_kind: String, + /// Whether the policy participates in new reservations. + pub enabled: bool, + /// Deterministic key of the period these amounts describe, absent when no + /// counter row exists yet for the current period. + pub period_key: Option, + /// Period start in Unix seconds, absent with `period_key`. + pub period_started_at: Option, + /// Period end in Unix seconds, absent with `period_key`. + pub period_ends_at: Option, + /// Input tokens currently held by unreconciled reservations. + pub reserved_input_tokens: i64, + /// Output tokens currently held by unreconciled reservations. + pub reserved_output_tokens: i64, + /// Tool units currently held by unreconciled reservations. + pub reserved_tool_units: i64, + /// Image units currently held by unreconciled reservations. + pub reserved_image_units: i64, + /// Cost microunits currently held by unreconciled reservations. + pub reserved_cost_microunits: i64, + /// Runs currently held by unreconciled reservations. + pub reserved_runs: i64, + /// Committed input tokens for the period. + pub committed_input_tokens: i64, + /// Committed output tokens for the period. + pub committed_output_tokens: i64, + /// Committed tool units for the period. + pub committed_tool_units: i64, + /// Committed image units for the period. + pub committed_image_units: i64, + /// Committed cost microunits for the period. + pub committed_cost_microunits: i64, + /// Committed runs for the period. + pub committed_runs: i64, + /// Configured input-token ceiling. + pub maximum_input_tokens: Option, + /// Configured output-token ceiling. + pub maximum_output_tokens: Option, + /// Configured tool-unit ceiling. + pub maximum_tool_units: Option, + /// Configured image-unit ceiling. + pub maximum_image_units: Option, + /// Configured cost ceiling in microunits. + pub maximum_cost_microunits: Option, + /// Configured run ceiling. + pub maximum_runs: Option, +} + +/// One unreconciled budget reservation visible to budget administration. +/// +/// It carries capacity accounting and durable run linkage only: never a +/// prompt, transcript, provider payload, principal identity, or credential. +#[derive(Clone, Debug, SimpleObject)] +#[cfg_attr(feature = "graphql-case-pascal", graphql(rename_fields = "PascalCase"))] +pub struct AiBudgetReservationCapacityView { + /// Reservation ID. + pub id: Uuid, + /// Owning run. + pub run_id: Uuid, + /// Stable reservation state: `reserved`, `uncertain`, or `committed` + /// after reclamation. + pub state: String, + /// Reservation expiry in Unix seconds. + pub expires_at: i64, + /// Creation time in Unix seconds. + pub created_at: i64, + /// Whether the reservation's expiry has passed. + pub expired: bool, + /// Whether the owning run reached a terminal state. + pub run_terminal: bool, + /// Whether the deployment is enabled and the durable time/run conditions + /// make this reservation a reclamation candidate. The mutation separately + /// rechecks current authorization, recent MFA, exact CAS, scope and stored + /// graph integrity; this field grants no authority and does not predict + /// those later checks. + pub reclaimable: bool, + /// Input tokens originally reserved. These count as held capacity only + /// while the reservation is unresolved. + pub reserved_input_tokens: i64, + /// Output tokens originally reserved. These count as held capacity only + /// while the reservation is unresolved. + pub reserved_output_tokens: i64, + /// Tool units originally reserved. These count as held capacity only + /// while the reservation is unresolved. + pub reserved_tool_units: i64, + /// Image units originally reserved. These count as held capacity only + /// while the reservation is unresolved. + pub reserved_image_units: i64, + /// Cost microunits originally reserved. These count as held capacity only + /// while the reservation is unresolved. + pub reserved_cost_microunits: i64, + /// Runs originally reserved. These count as held capacity only while the + /// reservation is unresolved. + pub reserved_runs: i64, + /// CAS version required to reclaim it. + pub row_version: i64, +} + +/// Bounded capacity and stranded-reservation report for one exact scope. +#[derive(Clone, Debug, SimpleObject)] +#[cfg_attr(feature = "graphql-case-pascal", graphql(rename_fields = "PascalCase"))] +pub struct AiBudgetScopeCapacityView { + /// Current-period capacity for every policy bound to the exact scope. + pub policies: Vec, + /// Reservations in state `uncertain` for the scope within the bounded + /// read window. + pub uncertain_reservation_count: i64, + /// Reservations in state `reserved` for the scope within the bounded read + /// window. + pub reserved_reservation_count: i64, + /// Unreconciled reservations whose expiry has already passed. + pub expired_reservation_count: i64, + /// Unreconciled reservations meeting the deployment and durable time/run + /// candidate conditions. Mutation authorization and CAS are separate. + pub reclaimable_reservation_count: i64, + /// Oldest unreconciled reservations first, bounded by the read window. + pub reservations: Vec, + /// Whether the bounded read window was filled, making every count above a + /// lower bound rather than an exact total. + pub truncated: bool, +} + +/// Exact one-shot privileged reclamation of one stranded reservation. +#[derive(Clone, Debug, InputObject)] +#[cfg_attr(feature = "graphql-case-pascal", graphql(rename_fields = "PascalCase"))] +pub struct ReclaimAiBudgetReservationInput { + /// Exact scope the reservation must belong to. + pub scope: AiScopeInput, + /// Reservation to resolve. + pub reservation_id: Uuid, + /// Expected reservation CAS version. + pub expected_version: i64, +} + /// Provider profile CAS upsert. #[derive(InputObject)] #[cfg_attr(feature = "graphql-case-pascal", graphql(rename_fields = "PascalCase"))] @@ -501,6 +655,53 @@ pub trait AiConfigurationService: Send + Sync { principal: &AuthPrincipal, input: UpsertAiBudgetPolicyInput, ) -> Result; + + /// Reports bounded reserved/committed capacity and stranded reservations + /// for one exact scope. + /// + /// The default implementation fails closed so an existing backend does not + /// gain a budget-observability surface implicitly. + /// + /// # Errors + /// + /// Returns [`AiError::InvalidConfiguration`] unless the backend implements + /// the surface; implementations return [`AiError::Forbidden`] for a denied + /// read and a safe library error for persistence failure. + async fn budget_scope_capacity( + &self, + _principal: &AuthPrincipal, + _scope: AiScope, + ) -> Result { + Err(AiError::InvalidConfiguration( + "budget capacity reporting is not implemented by this service".to_owned(), + )) + } + + /// Resolves one expired stranded reservation by committing the capacity it + /// already holds as authoritative usage, and audits the decision. + /// + /// This never proves that the provider was not reached. It commits the + /// held estimate, which can only over-count, so that capacity stops being + /// unreachable. The default implementation fails closed. + /// + /// # Errors + /// + /// Returns [`AiError::InvalidConfiguration`] unless the backend implements + /// and the deployment enables the surface, [`AiError::RecentMfaRequired`] + /// without current recent MFA, [`AiError::Forbidden`] for a denied action, + /// [`AiError::NotFound`] for an unknown reservation or one outside the + /// exact scope, and [`AiError::Conflict`] when the reservation is already + /// resolved, has not expired long enough, still has an active run, or + /// fails its CAS. + async fn reclaim_budget_reservation( + &self, + _principal: &AuthPrincipal, + _input: ReclaimAiBudgetReservationInput, + ) -> Result { + Err(AiError::InvalidConfiguration( + "budget reservation reclamation is not implemented by this service".to_owned(), + )) + } } /// Composable redacted configuration query root. @@ -579,6 +780,23 @@ impl AiConfigurationQueryRoot { Ok(policies) } + /// Reports bounded budget capacity and stranded reservations for a scope. + /// + /// Reserved capacity counts against the ceiling exactly like committed + /// usage, so a host should alarm on a rising unreconciled reservation + /// count before it becomes a refusal to serve. + async fn ai_budget_scope_capacity( + &self, + context: &Context<'_>, + scope: AiScopeInput, + ) -> async_graphql::Result { + let principal = agql_auth::principal_from_ctx(context)?; + configuration_service(context)? + .budget_scope_capacity(&principal, scope.into()) + .await + .map_err(extend) + } + /// Lists bounded immutable pricing versions for one exact route. async fn ai_pricing_policies( &self, @@ -697,6 +915,25 @@ impl AiConfigurationMutationRoot { .map_err(extend) } + /// Resolves one expired stranded budget reservation. + /// + /// The reserved amounts are committed as authoritative usage. That can + /// over-count and never under-counts, and it does not create headroom: the + /// reserved column falls by exactly the amount the committed column rises. + /// It exists so held capacity is accountable and reportable instead of + /// permanently unreachable. + async fn reclaim_ai_budget_reservation( + &self, + context: &Context<'_>, + input: ReclaimAiBudgetReservationInput, + ) -> async_graphql::Result { + let principal = agql_auth::principal_from_ctx(context)?; + configuration_service(context)? + .reclaim_budget_reservation(&principal, input) + .await + .map_err(extend) + } + /// Appends one immediately effective immutable pricing version. async fn create_ai_pricing_policy( &self, diff --git a/crates/graphql-orm-ai/src/orm_budget.rs b/crates/graphql-orm-ai/src/orm_budget.rs index 58523115..6e451c9f 100644 --- a/crates/graphql-orm-ai/src/orm_budget.rs +++ b/crates/graphql-orm-ai/src/orm_budget.rs @@ -684,10 +684,10 @@ impl AiBudgetService for OrmAiBudgetService { } #[derive(Clone, Debug)] -struct BudgetPeriod { - key: String, - started_at: i64, - ends_at: i64, +pub(crate) struct BudgetPeriod { + pub(crate) key: String, + pub(crate) started_at: i64, + pub(crate) ends_at: i64, } #[derive(Clone, Debug)] @@ -733,12 +733,30 @@ pub(crate) async fn commit_uncertain_background_budget( .ok_or_else(|| OrmPublicError::new(OrmErrorCode::InternalError))?; let counter_ids = reservation_counter_ids(record) .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + let reservation_scope = AiScope { + kind: record.scope_kind.clone(), + id: record.scope_id.clone(), + tenant_id: record.tenant_id.clone(), + }; for counter_id in counter_ids { let counter = tx .find_by_id::(&counter_id) .await .map_err(OrmPublicError::from)? .ok_or_else(OrmPublicError::not_found)?; + let policy = tx + .find_by_id::(&counter.budget_policy_id) + .await + .map_err(OrmPublicError::from)? + .ok_or_else(OrmPublicError::not_found)?; + if !policy_applies( + &policy, + &reservation_scope, + &record.principal_kind, + &record.principal_subject, + ) { + return Err(OrmPublicError::new(OrmErrorCode::InternalError)); + } let reserved_counter = counter_reserved(&counter) .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; let committed_counter = counter_committed(&counter) @@ -818,6 +836,147 @@ pub(crate) async fn commit_uncertain_background_budget( Ok(updated) } +/// Conservatively commits a stranded reservation's own reserved amounts as +/// authoritative usage inside the caller's wider state-machine transaction. +/// +/// This is the transactional half of the privileged reclamation surface. The +/// caller owns authorization, recent MFA, the exact scope binding, the expiry +/// and terminal-run evidence, and the audit append; this function owns the +/// counter arithmetic, usage insertion, and reservation CAS so reclamation +/// cannot diverge from ordinary provider reconciliation. +/// +/// It never invents an absence proof. A `Reserved` or `Uncertain` reservation +/// may have crossed the transport boundary, so the only safe resolution is to +/// charge the estimate that was already held. That can over-count, never +/// under-count, and it moves nothing between policies: the reserved column +/// falls by exactly the amount the committed column rises. +pub(crate) async fn commit_stranded_reservation( + tx: &mut graphql_orm::graphql::orm::MutationContext<'_, DefaultWriteBackend>, + record: &AiBudgetReservationRecord, + now: OffsetDateTime, +) -> Result { + let state = parse_reservation_state(&record.state) + .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + if !matches!( + state, + AiBudgetReservationState::Reserved | AiBudgetReservationState::Uncertain + ) { + return Err(OrmPublicError::new(OrmErrorCode::Conflict)); + } + let reserved = reservation_amounts(record) + .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + let actual = validate_reconciliation_actual( + reserved, + Some(reserved), + Some(0), + AiBudgetReconciliationOutcome::Commit, + ) + .map_err(|_| OrmPublicError::new(OrmErrorCode::Conflict))? + .ok_or_else(|| OrmPublicError::new(OrmErrorCode::InternalError))?; + let counter_ids = reservation_counter_ids(record) + .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + let reservation_scope = AiScope { + kind: record.scope_kind.clone(), + id: record.scope_id.clone(), + tenant_id: record.tenant_id.clone(), + }; + for counter_id in counter_ids { + let counter = tx + .find_by_id::(&counter_id) + .await + .map_err(OrmPublicError::from)? + .ok_or_else(OrmPublicError::not_found)?; + let policy = tx + .find_by_id::(&counter.budget_policy_id) + .await + .map_err(OrmPublicError::from)? + .ok_or_else(OrmPublicError::not_found)?; + if !policy_applies( + &policy, + &reservation_scope, + &record.principal_kind, + &record.principal_subject, + ) { + return Err(OrmPublicError::new(OrmErrorCode::InternalError)); + } + let reserved_counter = counter_reserved(&counter) + .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + let committed_counter = counter_committed(&counter) + .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + let next_reserved = checked_sub(reserved_counter, reserved) + .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + let next_committed = checked_add(committed_counter, actual) + .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + let update = counter_amount_update(next_reserved, next_committed)?; + if !matches!( + tx.compare_and_swap::( + &counter.id, + counter.row_version, + AiBudgetCounterRecordWhereInput::default(), + update, + ) + .await + .map_err(OrmPublicError::from)?, + ConditionalUpdateOutcome::Updated(_) + ) { + return Err(OrmPublicError::new(OrmErrorCode::Conflict)); + } + } + tx.insert::(CreateAiUsageEntryRecordInput { + id: Uuid::new_v4(), + budget_reservation_id: record.id, + scope_kind: record.scope_kind.clone(), + scope_id: record.scope_id.clone(), + tenant_id: record.tenant_id.clone(), + principal_kind: record.principal_kind.clone(), + principal_subject: record.principal_subject.clone(), + session_id: Some(record.session_id), + run_id: Some(record.run_id), + provider_kind: record.provider_kind.clone(), + provider_model: record.provider_model.clone(), + input_tokens: amount_to_i64(actual.input_tokens)?, + cached_input_tokens: 0, + output_tokens: amount_to_i64(actual.output_tokens)?, + tool_units: amount_to_i64(actual.tool_units)?, + image_units: amount_to_i64(actual.image_units)?, + cost_microunits: Some(amount_to_i64(actual.cost_microunits)?), + }) + .await + .map_err(OrmPublicError::from)?; + match tx + .compare_and_swap::( + &record.id, + record.row_version, + AiBudgetReservationRecordWhereInput { + state: Some(StringFilter { + eq: Some(record.state.clone()), + ..Default::default() + }), + ..Default::default() + }, + UpdateAiBudgetReservationRecordInput { + actual_input_tokens: Some(Some(amount_to_i64(actual.input_tokens)?)), + actual_cached_input_tokens: Some(Some(0)), + actual_output_tokens: Some(Some(amount_to_i64(actual.output_tokens)?)), + actual_tool_units: Some(Some(amount_to_i64(actual.tool_units)?)), + actual_image_units: Some(Some(amount_to_i64(actual.image_units)?)), + actual_cost_microunits: Some(Some(amount_to_i64(actual.cost_microunits)?)), + actual_runs: Some(Some(amount_to_i64(actual.runs)?)), + state: Some("committed".to_owned()), + reconciled_at: Some(Some(now.unix_timestamp())), + ..Default::default() + }, + ) + .await + .map_err(OrmPublicError::from)? + { + ConditionalUpdateOutcome::Updated(updated) => Ok(updated), + ConditionalUpdateOutcome::NotFound | ConditionalUpdateOutcome::Conflict => { + Err(OrmPublicError::new(OrmErrorCode::Conflict)) + } + } +} + #[cfg(feature = "provider-openai")] fn background_usage_identity(reservation_id: Uuid) -> Uuid { use sha2::{Digest, Sha256}; @@ -1035,7 +1194,7 @@ fn validate_policy_capacity( Ok(()) } -fn budget_period(interval: &str, now: OffsetDateTime) -> Result { +pub(crate) fn budget_period(interval: &str, now: OffsetDateTime) -> Result { let timestamp = now.unix_timestamp(); match interval { "minute" => fixed_period("minute", timestamp, 60), diff --git a/crates/graphql-orm-ai/src/orm_configuration.rs b/crates/graphql-orm-ai/src/orm_configuration.rs index e8ae18b4..7e2a8810 100644 --- a/crates/graphql-orm-ai/src/orm_configuration.rs +++ b/crates/graphql-orm-ai/src/orm_configuration.rs @@ -8,9 +8,10 @@ use agql_auth::{AuthPrincipal, Clock, RecentMfaPolicy}; use async_trait::async_trait; use graphql_orm::db::Database; use graphql_orm::graphql::errors::{OrmErrorCode, OrmPublicError}; -use graphql_orm::graphql::filters::StringFilter; +use graphql_orm::graphql::filters::{StringFilter, UuidFilter}; use graphql_orm::graphql::orm::{ - ConditionalUpdateOutcome, DefaultWriteBackend, TransactionError, TransactionMode, + ConditionalUpdateOutcome, DefaultWriteBackend, OrderDirection, TransactionError, + TransactionMode, }; use secrecy::SecretString; use serde::{Deserialize, Serialize}; @@ -18,13 +19,16 @@ use serde_json::json; use url::Url; use uuid::Uuid; +use crate::orm_budget::BudgetPeriod; use crate::persistence::*; use crate::{ - AiBudgetAmounts, AiBudgetPolicyView, AiConfigurationAccessPolicy, AiConfigurationAction, - AiConfigurationService, AiContentProtectionMode, AiContentProtectionPolicy, - AiContentProtectionPolicyResolver, AiContentProtectionPolicyView, AiError, - AiOpenAiCompatibleProfileInput, AiOpenAiCompatibleProfileView, AiProviderEndpointPolicy, - AiProviderKindInput, AiProviderProfileView, AiRetentionPolicyView, AiScope, AiSecretStore, + AiBudgetAmounts, AiBudgetPolicyCapacityView, AiBudgetPolicyView, + AiBudgetReservationCapacityView, AiBudgetScopeCapacityView, AiConfigurationAccessPolicy, + AiConfigurationAction, AiConfigurationService, AiContentProtectionMode, + AiContentProtectionPolicy, AiContentProtectionPolicyResolver, AiContentProtectionPolicyView, + AiError, AiOpenAiCompatibleProfileInput, AiOpenAiCompatibleProfileView, + AiProviderEndpointPolicy, AiProviderKindInput, AiProviderProfileView, AiRetentionPolicyView, + AiRunState, AiScope, AiSecretStore, ReclaimAiBudgetReservationInput, RemoveAiProviderCredentialInput, SecretRef, SetAiContentProtectionPolicyInput, SetAiRetentionPolicyInput, UpsertAiBudgetPolicyInput, UpsertAiProviderProfileInput, }; @@ -73,6 +77,74 @@ impl AiBudgetPolicyManagementLimits { } } +/// Deployment hard bounds for privileged budget-reservation reclamation. +/// +/// This type proves only that the deployment reviewed and enabled the surface +/// and chose how long an expired reservation must remain unresolved. It grants +/// no authority: the host still authorizes +/// [`AiConfigurationAction::ManageBudgetReclamation`] for the exact scope and +/// the caller still needs current recent MFA. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct AiBudgetReclamationLimits { + minimum_expired_age: time::Duration, + maximum_reservation_scan: usize, +} + +impl AiBudgetReclamationLimits { + /// Creates validated deployment reclamation bounds. + /// + /// `minimum_expired_age` is how long a reservation's `expires_at` must + /// already have passed before it may be resolved. It exists so an + /// in-flight provider turn whose worker is merely slow can never be + /// charged out from underneath itself. `maximum_reservation_scan` is a + /// deployment ceiling; the active database pagination maximum may narrow + /// the reported window further, in which case a full page is + /// conservatively marked truncated. + /// + /// # Errors + /// + /// Returns [`AiError::InvalidConfiguration`] unless the age is positive + /// and the bounded read window is in `1..=1000`. + pub fn new( + minimum_expired_age: time::Duration, + maximum_reservation_scan: usize, + ) -> Result { + if !minimum_expired_age.is_positive() || !(1..=1000).contains(&maximum_reservation_scan) { + return Err(AiError::InvalidConfiguration( + "invalid budget reclamation limits".to_owned(), + )); + } + Ok(Self { + minimum_expired_age, + maximum_reservation_scan, + }) + } + + /// Returns how long an expiry must already have passed. + pub const fn minimum_expired_age(self) -> time::Duration { + self.minimum_expired_age + } + + /// Returns the bounded reservation read window. + pub const fn maximum_reservation_scan(self) -> usize { + self.maximum_reservation_scan + } +} + +/// Bounded read window used when the deployment has not enabled reclamation. +const DEFAULT_BUDGET_RESERVATION_SCAN: usize = 200; + +/// Durable run states that can no longer reconcile their own reservation. +const TERMINAL_RUN_STATES: [&str; 4] = [ + AiRunState::Completed.as_str(), + AiRunState::Failed.as_str(), + AiRunState::Cancelled.as_str(), + AiRunState::RecoveryRequired.as_str(), +]; + +/// Reservation states that still hold capacity against a policy ceiling. +const UNRESOLVED_RESERVATION_STATES: [&str; 2] = ["reserved", "uncertain"]; + /// Durable configuration backend using generated ORM APIs and a compensating /// secret-reference saga. Secret plaintext never enters an ORM input. #[derive(Clone)] @@ -84,6 +156,7 @@ pub struct OrmAiConfigurationService { clock: Arc, secret_store: Arc, budget_policy_limits: Option, + budget_reclamation_limits: Option, } impl OrmAiConfigurationService { @@ -104,6 +177,7 @@ impl OrmAiConfigurationService { clock, secret_store, budget_policy_limits: None, + budget_reclamation_limits: None, } } @@ -116,6 +190,22 @@ impl OrmAiConfigurationService { self } + /// Enables privileged reclamation of stranded budget reservations under + /// deployment hard bounds. + /// + /// Without this explicit opt-in, capacity reporting still works and every + /// reservation reports `reclaimable: false`, while the mutation fails + /// closed as invalid configuration. Enabling it does not authorize anyone: + /// the host still decides + /// [`AiConfigurationAction::ManageBudgetReclamation`] per exact scope. + pub fn with_budget_reservation_reclamation( + mut self, + limits: AiBudgetReclamationLimits, + ) -> Self { + self.budget_reclamation_limits = Some(limits); + self + } + /// Returns the underlying ORM database handle for host schema wiring. pub fn database(&self) -> &Database { &self.database @@ -1004,6 +1094,315 @@ impl AiConfigurationService for OrmAiConfigurationService { .map_err(map_transaction)?; Ok(budget_policy_view(&record)) } + + async fn budget_scope_capacity( + &self, + principal: &AuthPrincipal, + scope: AiScope, + ) -> Result { + self.require_access(principal, &scope, AiConfigurationAction::ReadBudgetPolicies) + .await?; + let now = self.clock.now(); + let reclamation = self.budget_reclamation_limits; + let requested_window = reclamation.map_or(DEFAULT_BUDGET_RESERVATION_SCAN, |limits| { + limits.maximum_reservation_scan() + }); + // Typed ORM queries always honor the database's pagination ceiling. + // Narrow the administrative window to that ceiling instead of asking + // for a larger page that the ORM would silently clamp. When the + // ceiling leaves no room for a look-ahead record, a full page is + // conservatively reported as truncated. + let database_maximum = self + .database + .pagination_config() + .max_limit + .and_then(|maximum| usize::try_from(maximum.max(0)).ok()); + let window = + database_maximum.map_or(requested_window, |maximum| requested_window.min(maximum)); + let has_lookahead = database_maximum.is_none_or(|maximum| window < maximum); + let requested_scan = if has_lookahead { + window.saturating_add(1) + } else { + window + }; + let scan_limit = i64::try_from(requested_scan) + .map_err(|_| AiError::InvalidConfiguration("invalid scan window".to_owned()))?; + let exact_scope_key = scope_key(&scope); + let query_scope = scope.clone(); + let (policies, counters, reservations, run_reclamation_evidence) = self + .database + .transaction(TransactionMode::Default, move |tx| { + Box::pin(async move { + let policies = tx + .query::() + .filter(AiBudgetPolicyRecordWhereInput { + scope_key: Some(StringFilter { + eq: Some(exact_scope_key.clone()), + ..Default::default() + }), + ..Default::default() + }) + .default_order() + .limit(101) + .fetch_all() + .await + .map_err(OrmPublicError::from)?; + if policies.len() > 100 { + return Err(OrmPublicError::new(OrmErrorCode::PageLimitExceeded)); + } + if policies.iter().any(|record| { + record.scope_key != exact_scope_key + || record.scope_kind != query_scope.kind + || record.scope_id != query_scope.id + || record.tenant_id != query_scope.tenant_id + }) { + return Err(OrmPublicError::new(OrmErrorCode::InternalError)); + } + + let mut counters = Vec::with_capacity(policies.len()); + for policy in &policies { + let period = + crate::orm_budget::budget_period(&policy.interval_kind, now) + .map_err(|_| OrmPublicError::new(OrmErrorCode::InternalError))?; + let counter = tx + .query::() + .filter(AiBudgetCounterRecordWhereInput { + budget_policy_id: Some(UuidFilter { + eq: Some(policy.id), + ..Default::default() + }), + period_key: Some(StringFilter { + eq: Some(period.key.clone()), + ..Default::default() + }), + ..Default::default() + }) + .limit(1) + .fetch_one() + .await + .map_err(OrmPublicError::from)?; + counters.push((period, counter)); + } + + let tenant_id = Some(match &query_scope.tenant_id { + Some(tenant_id) => StringFilter { + eq: Some(tenant_id.clone()), + ..Default::default() + }, + None => StringFilter { + is_null: Some(true), + ..Default::default() + }, + }); + let reservations = tx + .query::() + .filter(AiBudgetReservationRecordWhereInput { + scope_kind: Some(StringFilter { + eq: Some(query_scope.kind.clone()), + ..Default::default() + }), + scope_id: Some(StringFilter { + eq: Some(query_scope.id.clone()), + ..Default::default() + }), + tenant_id, + state: Some(StringFilter { + in_list: Some( + UNRESOLVED_RESERVATION_STATES + .iter() + .map(|state| (*state).to_owned()) + .collect(), + ), + ..Default::default() + }), + ..Default::default() + }) + .order_by(AiBudgetReservationRecordOrderByInput { + created_at: Some(OrderDirection::Asc), + }) + .limit(scan_limit) + .fetch_all() + .await + .map_err(OrmPublicError::from)?; + if reservations + .iter() + .any(|record| record.tenant_id != query_scope.tenant_id) + { + return Err(OrmPublicError::new(OrmErrorCode::InternalError)); + } + + let mut run_reclamation_evidence = Vec::with_capacity(reservations.len()); + for reservation in &reservations { + let run = tx + .find_by_id::(&reservation.run_id) + .await + .map_err(OrmPublicError::from)?; + run_reclamation_evidence.push(run.map_or((false, false), |run| { + ( + run.session_id == reservation.session_id + && TERMINAL_RUN_STATES.contains(&run.state.as_str()), + run.session_id == reservation.session_id + && run.lease_owner.is_none() + && run.lease_expires_at.is_none(), + ) + })); + } + Ok((policies, counters, reservations, run_reclamation_evidence)) + }) + }) + .await + .map_err(map_transaction)?; + + let truncated = if has_lookahead { + reservations.len() > window + } else { + reservations.len() >= window + }; + let policy_views = policies + .iter() + .zip(counters.iter()) + .map(|(policy, (period, counter))| { + budget_policy_capacity_view(policy, period, counter.as_ref()) + }) + .collect::>(); + let mut uncertain_reservation_count = 0_i64; + let mut reserved_reservation_count = 0_i64; + let mut expired_reservation_count = 0_i64; + let mut reclaimable_reservation_count = 0_i64; + let mut reservation_views = Vec::with_capacity(reservations.len().min(window)); + for (record, (run_terminal, run_lease_free)) in reservations + .iter() + .take(window) + .zip(run_reclamation_evidence) + { + let view = budget_reservation_capacity_view( + record, + run_terminal, + run_lease_free, + now, + reclamation, + ); + if record.state == "uncertain" { + uncertain_reservation_count = uncertain_reservation_count.saturating_add(1); + } else { + reserved_reservation_count = reserved_reservation_count.saturating_add(1); + } + if view.expired { + expired_reservation_count = expired_reservation_count.saturating_add(1); + } + if view.reclaimable { + reclaimable_reservation_count = reclaimable_reservation_count.saturating_add(1); + } + reservation_views.push(view); + } + Ok(AiBudgetScopeCapacityView { + policies: policy_views, + uncertain_reservation_count, + reserved_reservation_count, + expired_reservation_count, + reclaimable_reservation_count, + reservations: reservation_views, + truncated, + }) + } + + async fn reclaim_budget_reservation( + &self, + principal: &AuthPrincipal, + input: ReclaimAiBudgetReservationInput, + ) -> Result { + self.require_recent_mfa(principal)?; + let scope: AiScope = input.scope.clone().into(); + self.require_access( + principal, + &scope, + AiConfigurationAction::ManageBudgetReclamation, + ) + .await?; + let limits = self.budget_reclamation_limits.ok_or_else(|| { + AiError::InvalidConfiguration( + "budget reservation reclamation is not enabled".to_owned(), + ) + })?; + if input.expected_version < 0 { + return Err(AiError::InvalidInput( + "invalid budget reservation version".to_owned(), + )); + } + let now = self.clock.now(); + let reclaimable_before = now + .checked_sub(limits.minimum_expired_age()) + .ok_or_else(|| AiError::InvalidConfiguration("budget time overflow".to_owned()))? + .unix_timestamp(); + let actor_kind = principal_kind(principal); + let actor_subject = principal.subject().to_owned(); + let reservation_id = input.reservation_id; + let expected_version = input.expected_version; + let updated = self + .database + .transaction(TransactionMode::StateMachine, move |tx| { + Box::pin(async move { + let record = tx + .find_by_id::(&reservation_id) + .await + .map_err(OrmPublicError::from)? + .ok_or_else(OrmPublicError::not_found)?; + if record.scope_kind != scope.kind + || record.scope_id != scope.id + || record.tenant_id != scope.tenant_id + { + return Err(OrmPublicError::not_found()); + } + if record.row_version != expected_version + || !UNRESOLVED_RESERVATION_STATES.contains(&record.state.as_str()) + || record.expires_at > reclaimable_before + { + return Err(OrmPublicError::new(OrmErrorCode::Conflict)); + } + // Authoritative evidence: the owning run reached a durable + // terminal state and holds no lease, so no worker can ever + // reconcile this reservation from its own transport + // knowledge. + let run = tx + .find_by_id::(&record.run_id) + .await + .map_err(OrmPublicError::from)? + .ok_or_else(|| OrmPublicError::new(OrmErrorCode::Conflict))?; + if run.session_id != record.session_id + || !TERMINAL_RUN_STATES.contains(&run.state.as_str()) + || run.lease_owner.is_some() + || run.lease_expires_at.is_some() + { + return Err(OrmPublicError::new(OrmErrorCode::Conflict)); + } + let previous_state = record.state.clone(); + let updated = + crate::orm_budget::commit_stranded_reservation(tx, &record, now).await?; + insert_audit( + tx, + AuditFact { + actor_principal_kind: &actor_kind, + actor_subject: &actor_subject, + action: "ai.budget_reservation.reclaim", + resource_kind: "budget_reservation", + resource_reference: &updated.id.to_string(), + outcome: "allowed", + reason_code: if previous_state == "uncertain" { + "expired_uncertain_reservation_committed" + } else { + "expired_reserved_reservation_committed" + }, + policy_version: Some(updated.row_version.to_string()), + }, + ) + .await?; + Ok(updated) + }) + }) + .await + .map_err(map_transaction)?; + Ok(budget_reclaimed_reservation_view(&updated)) + } } #[async_trait] @@ -1328,6 +1727,94 @@ fn budget_policy_view(record: &AiBudgetPolicyRecord) -> AiBudgetPolicyView { } } +fn budget_policy_capacity_view( + policy: &AiBudgetPolicyRecord, + period: &BudgetPeriod, + counter: Option<&AiBudgetCounterRecord>, +) -> AiBudgetPolicyCapacityView { + AiBudgetPolicyCapacityView { + policy_id: policy.id, + interval_kind: policy.interval_kind.clone(), + enabled: policy.enabled, + period_key: counter.map(|_| period.key.clone()), + period_started_at: counter.map(|_| period.started_at), + period_ends_at: counter.map(|_| period.ends_at), + reserved_input_tokens: counter.map_or(0, |row| row.reserved_input_tokens), + reserved_output_tokens: counter.map_or(0, |row| row.reserved_output_tokens), + reserved_tool_units: counter.map_or(0, |row| row.reserved_tool_units), + reserved_image_units: counter.map_or(0, |row| row.reserved_image_units), + reserved_cost_microunits: counter.map_or(0, |row| row.reserved_cost_microunits), + reserved_runs: counter.map_or(0, |row| row.reserved_runs), + committed_input_tokens: counter.map_or(0, |row| row.committed_input_tokens), + committed_output_tokens: counter.map_or(0, |row| row.committed_output_tokens), + committed_tool_units: counter.map_or(0, |row| row.committed_tool_units), + committed_image_units: counter.map_or(0, |row| row.committed_image_units), + committed_cost_microunits: counter.map_or(0, |row| row.committed_cost_microunits), + committed_runs: counter.map_or(0, |row| row.committed_runs), + maximum_input_tokens: policy.maximum_input_tokens, + maximum_output_tokens: policy.maximum_output_tokens, + maximum_tool_units: policy.maximum_tool_units, + maximum_image_units: policy.maximum_image_units, + maximum_cost_microunits: policy.maximum_cost_microunits, + maximum_runs: policy.maximum_runs, + } +} + +fn budget_reservation_capacity_view( + record: &AiBudgetReservationRecord, + run_terminal: bool, + run_lease_free: bool, + now: time::OffsetDateTime, + reclamation: Option, +) -> AiBudgetReservationCapacityView { + let expired = record.expires_at <= now.unix_timestamp(); + let reclaimable = run_terminal + && run_lease_free + && reclamation.is_some_and(|limits| { + now.checked_sub(limits.minimum_expired_age()) + .is_some_and(|threshold| record.expires_at <= threshold.unix_timestamp()) + }); + AiBudgetReservationCapacityView { + id: record.id, + run_id: record.run_id, + state: record.state.clone(), + expires_at: record.expires_at, + created_at: record.created_at, + expired, + run_terminal, + reclaimable, + reserved_input_tokens: record.reserved_input_tokens, + reserved_output_tokens: record.reserved_output_tokens, + reserved_tool_units: record.reserved_tool_units, + reserved_image_units: record.reserved_image_units, + reserved_cost_microunits: record.reserved_cost_microunits, + reserved_runs: record.reserved_runs, + row_version: record.row_version, + } +} + +fn budget_reclaimed_reservation_view( + record: &AiBudgetReservationRecord, +) -> AiBudgetReservationCapacityView { + AiBudgetReservationCapacityView { + id: record.id, + run_id: record.run_id, + state: record.state.clone(), + expires_at: record.expires_at, + created_at: record.created_at, + expired: true, + run_terminal: true, + reclaimable: false, + reserved_input_tokens: record.reserved_input_tokens, + reserved_output_tokens: record.reserved_output_tokens, + reserved_tool_units: record.reserved_tool_units, + reserved_image_units: record.reserved_image_units, + reserved_cost_microunits: record.reserved_cost_microunits, + reserved_runs: record.reserved_runs, + row_version: record.row_version, + } +} + fn exact_scope_key_for_record(record: &AiBudgetPolicyRecord) -> String { scope_key(&AiScope { kind: record.scope_kind.clone(), @@ -1539,3 +2026,840 @@ fn map_orm(error: OrmPublicError) -> AiError { | OrmErrorCode::AuthorizationMisconfigured => AiError::PersistenceFailed, } } + +#[cfg(all(test, feature = "sqlite"))] +mod tests { + use agql_auth::{ + AccessTokenMetadata, AssuranceMatchMode, AuthUser, FixedClock, MfaAcceptance, + ResolvedPrincipal, SessionAssurance, SessionContext, + }; + use graphql_orm::graphql::orm::{ApplyOptions, OrmSchemaModule}; + use graphql_orm::prelude::{Database, SqliteBackend}; + use time::Duration; + + use super::*; + use crate::orm_budget::{AiBudgetServiceLimits, OrmAiBudgetService}; + use crate::{ + AiBudgetReconciliation, AiBudgetReconciliationOutcome, AiBudgetReservationRequest, + AiBudgetService, AiRunId, AiSessionId, ModelReasoningEffort, ProviderKind, + }; + + const TENANT: &str = "tenant-1"; + const SUBJECT: &str = "budget-admin"; + const RESERVED_INPUT_TOKENS: u64 = 100; + + struct DenyConfiguration; + + #[async_trait] + impl AiConfigurationAccessPolicy for DenyConfiguration { + async fn can_configure( + &self, + _principal: &AuthPrincipal, + _scope: &AiScope, + _action: AiConfigurationAction, + ) -> bool { + false + } + } + + struct AllowConfiguration; + + #[async_trait] + impl AiConfigurationAccessPolicy for AllowConfiguration { + async fn can_configure( + &self, + _principal: &AuthPrincipal, + _scope: &AiScope, + _action: AiConfigurationAction, + ) -> bool { + true + } + } + + struct RejectEndpoints; + + impl AiProviderEndpointPolicy for RejectEndpoints { + fn authorize_endpoint( + &self, + _provider_kind: AiProviderKindInput, + _normalized_url: &str, + ) -> bool { + false + } + } + + #[derive(Default)] + struct UnusedSecretStore; + + #[async_trait] + impl AiSecretStore for UnusedSecretStore { + async fn resolve( + &self, + _reference: &SecretRef, + ) -> Result { + Err(crate::SecretError::Unavailable) + } + + async fn put( + &self, + _reference: Option<&SecretRef>, + _value: SecretString, + ) -> Result { + Err(crate::SecretError::Unavailable) + } + + async fn delete(&self, _reference: &SecretRef) -> Result<(), crate::SecretError> { + Ok(()) + } + } + + fn scope() -> AiScope { + AiScope::new("tenant", TENANT).with_tenant_id(TENANT) + } + + fn scope_input() -> crate::AiScopeInput { + crate::AiScopeInput { + kind: "tenant".to_owned(), + id: TENANT.to_owned(), + tenant_id: Some(TENANT.to_owned()), + } + } + + fn admin_principal(now: time::OffsetDateTime) -> AuthPrincipal { + let assurance = SessionAssurance::new( + now, + ["otp", "pwd"], + Some("urn:test:loa:2".to_owned()), + Some("test".to_owned()), + MfaAcceptance::Satisfied, + ) + .expect("test assurance should validate"); + AuthPrincipal::User(AuthUser { + user_id: SUBJECT.to_owned(), + session_id: Uuid::new_v4(), + roles: vec!["admin".to_owned()], + scopes: vec![], + session: SessionContext::default().with_assurance(assurance), + token_claims: AccessTokenMetadata { + auth_time: Some(now.unix_timestamp()), + amr: Some(vec!["otp".to_owned(), "pwd".to_owned()]), + acr: Some("urn:test:loa:2".to_owned()), + tenant_id: Some(TENANT.to_owned()), + ..AccessTokenMetadata::default() + }, + }) + } + + fn stale_mfa_principal() -> AuthPrincipal { + AuthPrincipal::User(AuthUser { + user_id: SUBJECT.to_owned(), + session_id: Uuid::new_v4(), + roles: vec!["admin".to_owned()], + scopes: vec![], + session: SessionContext::default(), + token_claims: AccessTokenMetadata { + tenant_id: Some(TENANT.to_owned()), + ..AccessTokenMetadata::default() + }, + }) + } + + fn configuration_service( + database: Database, + access_policy: Arc, + now: time::OffsetDateTime, + reclamation: bool, + ) -> OrmAiConfigurationService { + let service = OrmAiConfigurationService::new( + database, + access_policy, + Arc::new(RejectEndpoints), + RecentMfaPolicy { + maximum_age: Duration::minutes(5), + clock_skew: Duration::seconds(30), + allowed_amr: vec!["otp".to_owned()], + allowed_acr: vec!["urn:test:loa:2".to_owned()], + match_mode: AssuranceMatchMode::All, + }, + Arc::new(FixedClock::new(now)), + Arc::new(UnusedSecretStore), + ); + if reclamation { + service.with_budget_reservation_reclamation( + AiBudgetReclamationLimits::new(Duration::hours(1), 200) + .expect("reclamation limits should validate"), + ) + } else { + service + } + } + + /// Seeds a policy, session, running run, and one `uncertain` reservation + /// created through the real budget service and marked uncertain through the + /// real transport-boundary reconciliation. + async fn stranded_reservation_fixture() -> (Database, time::OffsetDateTime, Uuid) + { + let database = Database::::connect_sqlite("sqlite::memory:") + .await + .expect("in-memory SQLite should open"); + let module = crate::AiSchemaModule; + let plan = database + .schema() + .plan_migration_to_entities( + "ai-budget-reclaim-test-v1", + "AI budget reclamation test", + module.entities(), + ) + .await + .expect("AI schema migration should plan"); + database + .schema() + .apply_migration(&plan, ApplyOptions::default()) + .await + .expect("AI schema migration should apply"); + + let now = time::OffsetDateTime::from_unix_timestamp(1_800_000_000) + .expect("fixed test timestamp should be valid"); + AiBudgetPolicyRecord::insert( + &database, + CreateAiBudgetPolicyRecordInput { + scope_key: scope_key(&scope()), + scope_kind: "tenant".to_owned(), + scope_id: TENANT.to_owned(), + tenant_id: Some(TENANT.to_owned()), + principal_kind: None, + principal_subject: None, + interval_kind: "month".to_owned(), + maximum_input_tokens: Some(1_000), + maximum_output_tokens: Some(1_000), + maximum_tool_units: Some(100), + maximum_image_units: Some(100), + maximum_cost_microunits: Some(10_000), + maximum_runs: Some(100), + enabled: true, + }, + ) + .await + .expect("budget policy should seed"); + + let principal = admin_principal(now); + let resolved = ResolvedPrincipal::new(principal.reference(), principal.clone(), now) + .expect("fresh principal should resolve"); + let session_id = AiSessionId::new(); + let run_id = AiRunId::new(); + let attempt_id = Uuid::new_v4(); + AiSessionRecord::insert( + &database, + CreateAiSessionRecordInput { + id: session_id.0, + owner_principal_kind: "user".to_owned(), + owner_subject: SUBJECT.to_owned(), + tenant_id: Some(TENANT.to_owned()), + scope_kind: "tenant".to_owned(), + scope_id: TENANT.to_owned(), + title: "Budget reclamation".to_owned(), + title_revision: 0, + title_source: "default".to_owned(), + state: "active".to_owned(), + stream_head: 0, + message_head: 0, + last_activity_at: now.unix_timestamp(), + archived_at: None, + deleted_at: None, + }, + ) + .await + .expect("session should seed"); + AiRunRecord::insert( + &database, + CreateAiRunRecordInput { + id: run_id.0, + session_id: session_id.0, + input_message_id: Uuid::new_v4(), + principal_reference: serde_json::to_value(resolved.reference()) + .expect("principal reference should serialize"), + state: "running".to_owned(), + attempt_id: Some(attempt_id), + lease_owner: Some("worker-test".to_owned()), + lease_generation: 1, + lease_expires_at: Some((now + Duration::minutes(4)).unix_timestamp()), + lease_heartbeat_at: Some(now.unix_timestamp()), + retry_count: 0, + next_attempt_at: None, + error_code: None, + latest_checkpoint_id: None, + cancellation_request_id: None, + cancellation_requested_at: None, + }, + ) + .await + .expect("running run should seed"); + + let budget = OrmAiBudgetService::new( + database.clone(), + Arc::new(FixedClock::new(now)), + AiBudgetServiceLimits::new( + AiBudgetAmounts { + input_tokens: 1_000, + output_tokens: 1_000, + tool_units: 100, + image_units: 100, + cost_microunits: 10_000, + runs: 1, + }, + Duration::minutes(5), + Duration::seconds(30), + 16, + 8, + ) + .expect("budget service limits should validate"), + ); + let reservation = budget + .reserve( + &resolved, + AiBudgetReservationRequest { + scope: scope(), + session_id, + run_id, + attempt_id, + lease_generation: 1, + provider_kind: ProviderKind::OpenAi, + model: "test-model".to_owned(), + reasoning_effort: ModelReasoningEffort::Unspecified, + pricing_policy_version: "pricing-test-v1".to_owned(), + estimate: AiBudgetAmounts { + input_tokens: RESERVED_INPUT_TOKENS, + output_tokens: 10, + tool_units: 0, + image_units: 0, + cost_microunits: 100, + runs: 1, + }, + idempotency_key: "reclaim-test-1".to_owned(), + expires_at: now + Duration::minutes(2), + }, + ) + .await + .expect("reservation should be granted"); + budget + .reconcile( + &resolved, + AiBudgetReconciliation { + reservation_id: reservation.id(), + attempt_id, + lease_generation: 1, + actual: None, + cached_input_tokens: None, + outcome: AiBudgetReconciliationOutcome::MarkUncertain, + }, + ) + .await + .expect("transport boundary should mark the reservation uncertain"); + (database, now, reservation.id().0) + } + + async fn terminate_run(database: &Database) { + let run = database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .limit(2) + .fetch_all() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("run query should succeed") + .into_iter() + .next() + .expect("one run was seeded"); + AiRunRecord::update_by_id( + database, + &run.id, + UpdateAiRunRecordInput { + state: Some("recovery_required".to_owned()), + attempt_id: Some(None), + lease_owner: Some(None), + lease_expires_at: Some(None), + lease_heartbeat_at: Some(None), + error_code: Some(Some("provider_turn_uncertain".to_owned())), + ..Default::default() + }, + ) + .await + .expect("run should reach a terminal state"); + } + + async fn make_run_terminal_without_releasing_lease(database: &Database) { + let run = database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .limit(2) + .fetch_all() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("run query should succeed") + .into_iter() + .next() + .expect("one run was seeded"); + AiRunRecord::update_by_id( + database, + &run.id, + UpdateAiRunRecordInput { + state: Some("recovery_required".to_owned()), + error_code: Some(Some("provider_turn_uncertain".to_owned())), + ..Default::default() + }, + ) + .await + .expect("test should create inconsistent terminal lease evidence"); + } + + async fn expire_reservation( + database: &Database, + reservation_id: Uuid, + expires_at: i64, + ) { + AiBudgetReservationRecord::update_by_id( + database, + &reservation_id, + UpdateAiBudgetReservationRecordInput { + expires_at: Some(expires_at), + ..Default::default() + }, + ) + .await + .expect("reservation expiry should rewind"); + } + + async fn usage_entries(database: &Database) -> Vec { + database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .limit(10) + .fetch_all() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("usage query should succeed") + } + + async fn audit_actions(database: &Database) -> Vec { + database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .limit(20) + .fetch_all() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("audit query should succeed") + .into_iter() + .map(|record| record.action) + .collect() + } + + async fn seed_other_tenant_reservation(database: &Database) -> Uuid { + let original = database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .limit(2) + .fetch_all() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("reservation query should succeed") + .into_iter() + .next() + .expect("one reservation was seeded"); + let inserted = AiBudgetReservationRecord::insert( + database, + CreateAiBudgetReservationRecordInput { + budget_counter_ids: original.budget_counter_ids, + scope_kind: original.scope_kind, + scope_id: original.scope_id, + tenant_id: Some("other-tenant".to_owned()), + principal_kind: original.principal_kind, + principal_subject: "other-user".to_owned(), + session_id: original.session_id, + run_id: original.run_id, + attempt_id: original.attempt_id, + lease_generation: original.lease_generation, + provider_kind: original.provider_kind, + provider_model: original.provider_model, + reasoning_effort: original.reasoning_effort, + pricing_policy_version: original.pricing_policy_version, + reserved_input_tokens: original.reserved_input_tokens, + reserved_output_tokens: original.reserved_output_tokens, + reserved_tool_units: original.reserved_tool_units, + reserved_image_units: original.reserved_image_units, + reserved_cost_microunits: original.reserved_cost_microunits, + reserved_runs: original.reserved_runs, + actual_input_tokens: original.actual_input_tokens, + actual_cached_input_tokens: original.actual_cached_input_tokens, + actual_output_tokens: original.actual_output_tokens, + actual_tool_units: original.actual_tool_units, + actual_image_units: original.actual_image_units, + actual_cost_microunits: original.actual_cost_microunits, + actual_runs: original.actual_runs, + idempotency_key: "other-tenant-reservation".to_owned(), + state: original.state, + expires_at: original.expires_at, + reconciled_at: original.reconciled_at, + }, + ) + .await + .expect("other-tenant reservation should seed"); + inserted.id + } + + #[tokio::test] + async fn expired_uncertain_reservation_is_reclaimable_and_reported() { + let (database, now, reservation_id) = stranded_reservation_fixture().await; + let service = + configuration_service(database.clone(), Arc::new(AllowConfiguration), now, true); + let principal = admin_principal(now); + + let before = service + .budget_scope_capacity(&principal, scope()) + .await + .expect("capacity reporting should succeed"); + assert_eq!(before.policies.len(), 1); + assert_eq!( + before.policies[0].reserved_input_tokens, + RESERVED_INPUT_TOKENS as i64 + ); + assert_eq!(before.policies[0].committed_input_tokens, 0); + assert_eq!(before.uncertain_reservation_count, 1); + assert_eq!(before.reserved_reservation_count, 0); + assert!(!before.truncated); + assert_eq!(before.reservations.len(), 1); + assert!(!before.reservations[0].expired); + assert!(!before.reservations[0].run_terminal); + assert!(!before.reservations[0].reclaimable); + + terminate_run(&database).await; + expire_reservation( + &database, + reservation_id, + (now - Duration::days(2)).unix_timestamp(), + ) + .await; + + let stranded = service + .budget_scope_capacity(&principal, scope()) + .await + .expect("capacity reporting should succeed"); + assert_eq!(stranded.expired_reservation_count, 1); + assert_eq!(stranded.reclaimable_reservation_count, 1); + assert!(stranded.reservations[0].reclaimable); + let expected_version = stranded.reservations[0].row_version; + + assert!(matches!( + service + .reclaim_budget_reservation( + &principal, + ReclaimAiBudgetReservationInput { + scope: scope_input(), + reservation_id, + expected_version: expected_version + 5, + }, + ) + .await, + Err(AiError::Conflict) + )); + + let reclaimed = service + .reclaim_budget_reservation( + &principal, + ReclaimAiBudgetReservationInput { + scope: scope_input(), + reservation_id, + expected_version, + }, + ) + .await + .expect("an expired uncertain reservation on a terminal run reclaims"); + assert_eq!(reclaimed.state, "committed"); + assert_eq!( + reclaimed.reserved_input_tokens, + RESERVED_INPUT_TOKENS as i64 + ); + + let after = service + .budget_scope_capacity(&principal, scope()) + .await + .expect("capacity reporting should succeed"); + assert_eq!(after.policies[0].reserved_input_tokens, 0); + assert_eq!( + after.policies[0].committed_input_tokens, + RESERVED_INPUT_TOKENS as i64 + ); + assert_eq!(after.uncertain_reservation_count, 0); + assert_eq!(after.reclaimable_reservation_count, 0); + assert!(after.reservations.is_empty()); + + let usage = usage_entries(&database).await; + assert_eq!(usage.len(), 1); + assert_eq!(usage[0].budget_reservation_id, reservation_id); + assert_eq!(usage[0].input_tokens, RESERVED_INPUT_TOKENS as i64); + assert_eq!(usage[0].cached_input_tokens, 0); + assert!( + audit_actions(&database) + .await + .contains(&"ai.budget_reservation.reclaim".to_owned()) + ); + + // Replay of the same exact-version request cannot double-count. + assert!(matches!( + service + .reclaim_budget_reservation( + &principal, + ReclaimAiBudgetReservationInput { + scope: scope_input(), + reservation_id, + expected_version, + }, + ) + .await, + Err(AiError::Conflict) + )); + } + + #[tokio::test] + async fn unexpired_or_active_uncertain_reservations_are_not_reclaimable() { + let (database, now, reservation_id) = stranded_reservation_fixture().await; + let service = + configuration_service(database.clone(), Arc::new(AllowConfiguration), now, true); + let principal = admin_principal(now); + let version = service + .budget_scope_capacity(&principal, scope()) + .await + .expect("capacity reporting should succeed") + .reservations[0] + .row_version; + + // Expired long ago, but the owning run is still running. + expire_reservation( + &database, + reservation_id, + (now - Duration::days(2)).unix_timestamp(), + ) + .await; + assert!(matches!( + service + .reclaim_budget_reservation( + &principal, + ReclaimAiBudgetReservationInput { + scope: scope_input(), + reservation_id, + expected_version: version, + }, + ) + .await, + Err(AiError::Conflict) + )); + + // A terminal state alone is insufficient while stale lease evidence + // remains. Reporting and mutation must agree on the same fail-closed + // reclaimability predicate. + make_run_terminal_without_releasing_lease(&database).await; + let terminal_but_leased = service + .budget_scope_capacity(&principal, scope()) + .await + .expect("capacity reporting should succeed"); + assert!(terminal_but_leased.reservations[0].run_terminal); + assert!(!terminal_but_leased.reservations[0].reclaimable); + assert!(matches!( + service + .reclaim_budget_reservation( + &principal, + ReclaimAiBudgetReservationInput { + scope: scope_input(), + reservation_id, + expected_version: terminal_but_leased.reservations[0].row_version, + }, + ) + .await, + Err(AiError::Conflict) + )); + + // Terminal run, but the expiry grace has not elapsed. + terminate_run(&database).await; + expire_reservation( + &database, + reservation_id, + (now - Duration::minutes(1)).unix_timestamp(), + ) + .await; + let capacity = service + .budget_scope_capacity(&principal, scope()) + .await + .expect("capacity reporting should succeed"); + assert!(capacity.reservations[0].expired); + assert!(capacity.reservations[0].run_terminal); + assert!(!capacity.reservations[0].reclaimable); + assert!(matches!( + service + .reclaim_budget_reservation( + &principal, + ReclaimAiBudgetReservationInput { + scope: scope_input(), + reservation_id, + expected_version: capacity.reservations[0].row_version, + }, + ) + .await, + Err(AiError::Conflict) + )); + + // Nothing moved between the reserved and committed columns. + assert_eq!( + capacity.policies[0].reserved_input_tokens, + RESERVED_INPUT_TOKENS as i64 + ); + assert_eq!(capacity.policies[0].committed_input_tokens, 0); + assert!(usage_entries(&database).await.is_empty()); + } + + #[tokio::test] + async fn reclamation_requires_authorization_recent_mfa_and_deployment_opt_in() { + let (database, now, reservation_id) = stranded_reservation_fixture().await; + terminate_run(&database).await; + expire_reservation( + &database, + reservation_id, + (now - Duration::days(2)).unix_timestamp(), + ) + .await; + let principal = admin_principal(now); + let input = || ReclaimAiBudgetReservationInput { + scope: scope_input(), + reservation_id, + expected_version: 1, + }; + + let denied = + configuration_service(database.clone(), Arc::new(DenyConfiguration), now, true); + assert!(matches!( + denied.reclaim_budget_reservation(&principal, input()).await, + Err(AiError::Forbidden) + )); + assert!(matches!( + denied.budget_scope_capacity(&principal, scope()).await, + Err(AiError::Forbidden) + )); + + let allowed = + configuration_service(database.clone(), Arc::new(AllowConfiguration), now, true); + assert!(matches!( + allowed + .reclaim_budget_reservation(&stale_mfa_principal(), input()) + .await, + Err(AiError::RecentMfaRequired) + )); + + let unconfigured = + configuration_service(database.clone(), Arc::new(AllowConfiguration), now, false); + assert!(matches!( + unconfigured + .reclaim_budget_reservation(&principal, input()) + .await, + Err(AiError::InvalidConfiguration(_)) + )); + assert!( + !unconfigured + .budget_scope_capacity(&principal, scope()) + .await + .expect("capacity reporting works without the reclamation opt-in") + .reservations[0] + .reclaimable + ); + + // No refused path may move capacity or append usage. + assert!(usage_entries(&database).await.is_empty()); + } + + #[tokio::test] + async fn capacity_reporting_filters_the_exact_tenant_before_its_bound() { + let (database, now, reservation_id) = stranded_reservation_fixture().await; + let _other_reservation = seed_other_tenant_reservation(&database).await; + let service = configuration_service(database, Arc::new(AllowConfiguration), now, true); + + let capacity = service + .budget_scope_capacity(&admin_principal(now), scope()) + .await + .expect("exact-tenant capacity reporting should succeed"); + + assert_eq!(capacity.reservations.len(), 1); + assert_eq!(capacity.reservations[0].id, reservation_id); + assert!(!capacity.truncated); + } + + #[tokio::test] + async fn reclamation_rejects_a_cross_scope_counter_link() { + let (database, now, _reservation_id) = stranded_reservation_fixture().await; + let other_reservation_id = seed_other_tenant_reservation(&database).await; + terminate_run(&database).await; + expire_reservation( + &database, + other_reservation_id, + (now - Duration::days(2)).unix_timestamp(), + ) + .await; + let service = + configuration_service(database.clone(), Arc::new(AllowConfiguration), now, true); + let other_scope = AiScope::new("tenant", TENANT).with_tenant_id("other-tenant"); + let candidate = service + .budget_scope_capacity(&admin_principal(now), other_scope.clone()) + .await + .expect("the malformed reservation remains observable") + .reservations + .into_iter() + .next() + .expect("the other-tenant reservation should be visible"); + assert!(candidate.reclaimable); + + assert!(matches!( + service + .reclaim_budget_reservation( + &admin_principal(now), + ReclaimAiBudgetReservationInput { + scope: crate::AiScopeInput { + kind: other_scope.kind, + id: other_scope.id, + tenant_id: other_scope.tenant_id, + }, + reservation_id: other_reservation_id, + expected_version: candidate.row_version, + }, + ) + .await, + Err(AiError::PersistenceFailed) + )); + assert!(usage_entries(&database).await.is_empty()); + let original_capacity = service + .budget_scope_capacity(&admin_principal(now), scope()) + .await + .expect("the original scope capacity remains readable"); + assert_eq!( + original_capacity.policies[0].reserved_input_tokens, + RESERVED_INPUT_TOKENS as i64 + ); + assert_eq!(original_capacity.policies[0].committed_input_tokens, 0); + } +} diff --git a/crates/graphql-orm-ai/src/orm_coordinator.rs b/crates/graphql-orm-ai/src/orm_coordinator.rs index 6a1ae2e5..78c63998 100644 --- a/crates/graphql-orm-ai/src/orm_coordinator.rs +++ b/crates/graphql-orm-ai/src/orm_coordinator.rs @@ -408,6 +408,13 @@ impl AiAgentRunControl for OrmAiRunService { pub trait AiAgentProviderTurnExecutor: Send + Sync { /// Executes one exactly planned turn for the current attempt/generation. /// + /// Returning [`AiError::PreTransportBudgetDenied`] is a proof-bearing + /// contract: the + /// denial must have occurred before provider transport, and no budget + /// reservation may remain held. Once transport might have occurred, an + /// executor must return [`AiError::ProviderFailed`] instead so the + /// coordinator preserves uncertainty. + /// /// # Errors /// /// Returns a safe library error for any authorization, budget, egress, @@ -1506,6 +1513,19 @@ impl AiReadOnlyAgentCoordinator { ) .await; } + Err(ProviderTurnFailure::BudgetDenied) => { + // The atomic reservation is taken before the transport + // boundary, so a denial is a certain, local, pre-transport + // refusal: no provider turn was consumed and the + // reservation transaction held nothing. Closing the run + // for recovery here would report a proven refusal as + // unprovable provider uncertainty and permanently refuse + // retry admission for a run that is safe to author again + // once capacity exists. + return self + .finish_failed(&lease, &guard, "provider_budget_denied") + .await; + } Err(ProviderTurnFailure::LeaseLost(error)) => return Err(error), Err(ProviderTurnFailure::Cancelled(settlement)) => { self.settle_interrupted_provider_session(&lease, &guard, settlement) @@ -1932,7 +1952,9 @@ impl AiReadOnlyAgentCoordinator { .wait_for_cancellation(&cancellation_lease, heartbeat_delay); tokio::pin!(cancellation); tokio::select! { - result = &mut provider => return result.map_err(|_| ProviderTurnFailure::Provider), + result = &mut provider => { + return result.map_err(|error| classify_provider_turn_failure(&error)); + } result = &mut cancellation => { match result.map_err(ProviderTurnFailure::LeaseLost)? { Some(_) => { @@ -1979,7 +2001,7 @@ impl AiReadOnlyAgentCoordinator { tokio::select! { result = &mut provider => { *lease = lease_state.lock().await.clone(); - return result.map_err(|_| ProviderTurnFailure::Provider); + return result.map_err(|error| classify_provider_turn_failure(&error)); } result = &mut cancellation => { match result.map_err(ProviderTurnFailure::LeaseLost)? { @@ -2049,7 +2071,7 @@ impl AiReadOnlyAgentCoordinator { Err(AiError::ProviderSessionDeferred) => { Err(ProviderTurnFailure::Deferred) } - Err(_) => Err(ProviderTurnFailure::Provider), + Err(error) => Err(classify_provider_turn_failure(&error)), }; } result = &mut cancellation => { @@ -2251,6 +2273,7 @@ impl AiReadOnlyAgentCoordinator { enum ProviderTurnFailure { Provider, + BudgetDenied, Deferred, LeaseLost(AiError), /// Owner cancellation won the fence; the value reports what the resulting @@ -2258,6 +2281,20 @@ enum ProviderTurnFailure { Cancelled(crate::AiRunInterruptSettlement), } +/// Separates a certain pre-transport refusal from an uncertain provider turn. +/// +/// The budget reservation is taken before the transport boundary and inside +/// the same call that later dispatches. A denial therefore proves that no +/// bytes crossed the provider boundary, that no provider turn was consumed, +/// and that the atomic reservation transaction left nothing held. Every other +/// executor error keeps the fail-closed uncertain classification. +const fn classify_provider_turn_failure(error: &AiError) -> ProviderTurnFailure { + match error { + AiError::PreTransportBudgetDenied => ProviderTurnFailure::BudgetDenied, + _ => ProviderTurnFailure::Provider, + } +} + #[cfg(all(test, feature = "sqlite"))] mod tests { use std::collections::VecDeque; @@ -3554,6 +3591,108 @@ mod tests { )); } + #[tokio::test] + async fn pre_transport_budget_denial_fails_cleanly_instead_of_requiring_recovery() { + let lease = AiRunLease::test_running(principal_reference()); + let run = Arc::new(TestRunControl::new()); + let provider = Arc::new(TestProviderExecutor { + responses: Mutex::new(VecDeque::from([Err(AiError::PreTransportBudgetDenied)])), + delay: None, + }); + let planner = Arc::new(TestChatPlanner { + scope: test_scope(), + continuation_count: AtomicUsize::new(0), + }); + let forbidden = Arc::new(ChatForbiddenBoundaries::default()); + let coordinator = AiReadOnlyAgentCoordinator::new( + run.clone(), + provider.clone(), + forbidden.clone(), + Arc::new(TestOutputWriter), + forbidden.clone(), + Arc::new(TestCheckpointWriter), + Arc::new(TestRuleResolver), + planner, + limits(50), + ); + + let outcome = coordinator + .execute_claimed(&lease) + .await + .expect("a budget denial is a clean terminal failure"); + + assert!(matches!( + outcome, + Failed { + provider_turns: 0, + total_tool_calls: 0, + } + )); + assert_eq!(run.final_states(), vec![AiRunState::Failed]); + assert_eq!(run.final_codes(), vec!["provider_budget_denied".to_owned()]); + assert!(run.scheduled_retry_codes().is_empty()); + assert_eq!(provider.remaining_responses(), 0); + assert_eq!(forbidden.tool_calls.load(Ordering::SeqCst), 0); + assert_eq!(forbidden.provider_checkpoints.load(Ordering::SeqCst), 0); + // A certain pre-transport refusal is safe to author again once + // capacity exists, so the terminal-event failure record admits retry. + assert_eq!( + crate::classify_run_retry( + crate::AiRunRetryEvidence { + terminal: crate::AiRunTerminalEvent::Failed, + produced_assistant_output: false, + }, + Some("provider_budget_denied"), + ), + crate::AiRunRetryAdmission::Allowed + ); + } + + #[tokio::test] + async fn generic_budget_denial_cannot_claim_pre_transport_certainty() { + let lease = AiRunLease::test_running(principal_reference()); + let run = Arc::new(TestRunControl::new()); + let provider = Arc::new(TestProviderExecutor { + responses: Mutex::new(VecDeque::from([Err(AiError::BudgetDenied)])), + delay: None, + }); + let planner = Arc::new(TestChatPlanner { + scope: test_scope(), + continuation_count: AtomicUsize::new(0), + }); + let forbidden = Arc::new(ChatForbiddenBoundaries::default()); + let coordinator = AiReadOnlyAgentCoordinator::new( + run.clone(), + provider, + forbidden.clone(), + Arc::new(TestOutputWriter), + forbidden.clone(), + Arc::new(TestCheckpointWriter), + Arc::new(TestRuleResolver), + planner, + limits(50), + ); + + let outcome = coordinator + .execute_claimed(&lease) + .await + .expect("a generic denial must preserve possible transport uncertainty"); + + assert!(matches!( + outcome, + RecoveryRequired { + phase: AiAgentRecoveryPhase::ProviderTurn, + provider_turns: 0, + total_tool_calls: 0, + } + )); + assert_eq!(run.final_states(), vec![AiRunState::RecoveryRequired]); + assert_eq!( + run.final_codes(), + vec!["provider_turn_uncertain".to_owned()] + ); + } + #[tokio::test] async fn chat_turn_persists_final_output_without_tool_boundaries() { let lease = AiRunLease::test_running(principal_reference()); @@ -4197,6 +4336,71 @@ mod tests { assert_eq!(run.final_states(), vec![AiRunState::Completed]); } + #[tokio::test] + async fn post_transport_dynamic_tool_limit_remains_provider_uncertainty() { + let lease = AiRunLease::test_running(principal_reference()); + let run = Arc::new(TestRunControl::new()); + let provider = Arc::new(TestProviderExecutor { + responses: Mutex::new(VecDeque::from([Ok(AiProviderCallResult::test_result( + &lease, + None, + "response-dynamic-tool-limit", + vec![ + ("dynamic-call-1", "test.read", json!({})), + ("dynamic-call-2", "test.read", json!({})), + ], + ))])), + delay: None, + }); + let planner = Arc::new(TestDynamicPlanner { + scope: test_scope(), + route: test_route(), + continuation_count: AtomicUsize::new(0), + }); + let coordinator_limits = AiReadOnlyAgentCoordinatorLimits::new( + AiAgentLoopLimits::new(4, 1).expect("test loop limits should validate"), + Duration::milliseconds(50), + ) + .expect("test coordinator limits should validate"); + let coordinator = AiReadOnlyAgentCoordinator::new( + run.clone(), + provider, + Arc::new(TestToolExecutor { + expose_result: true, + }), + Arc::new(TestOutputWriter), + Arc::new(ChatForbiddenBoundaries::default()), + Arc::new(TestCheckpointWriter), + Arc::new(TestRuleResolver), + planner, + coordinator_limits, + ); + + let outcome = coordinator + .execute_claimed(&lease) + .await + .expect("a limit reached after dispatch must retain uncertainty"); + + assert!(matches!( + outcome, + RecoveryRequired { + phase: AiAgentRecoveryPhase::ProviderTurn, + provider_turns: 0, + total_tool_calls: 0, + } + )); + assert_eq!(run.final_states(), vec![AiRunState::RecoveryRequired]); + assert_eq!( + run.final_codes(), + vec!["provider_turn_uncertain".to_owned()] + ); + assert!( + !run.final_codes() + .iter() + .any(|code| code == "provider_budget_denied") + ); + } + #[tokio::test] async fn experimental_dynamic_turn_denies_rule_change_before_tool_execution() { let lease = AiRunLease::test_running(principal_reference()); diff --git a/crates/graphql-orm-ai/src/orm_supervised_coordinator.rs b/crates/graphql-orm-ai/src/orm_supervised_coordinator.rs index b02c84ef..027469f2 100644 --- a/crates/graphql-orm-ai/src/orm_supervised_coordinator.rs +++ b/crates/graphql-orm-ai/src/orm_supervised_coordinator.rs @@ -1087,6 +1087,18 @@ impl AiSupervisedAgentCoordinator { ) .await; } + Err(SupervisedProviderTurnFailure::BudgetDenied) => { + if let Some((service, claim)) = &reclaimed { + let _ = service + .require_cleanup(claim, "provider_session_reclaimed_handoff_failed") + .await; + } + // A pre-transport reservation denial is certain and local: it + // consumed no provider turn and left no reservation held. + return self + .finish_failed(&lease, &guard, "provider_budget_denied") + .await; + } Err(SupervisedProviderTurnFailure::LeaseLost(error)) => return Err(error), }; if self.run_control.cancellation(&lease).await?.is_some() { @@ -1588,7 +1600,7 @@ impl AiSupervisedAgentCoordinator { tokio::pin!(heartbeat); tokio::select! { result = &mut provider => { - return result.map_err(|_| SupervisedProviderTurnFailure::Provider); + return result.map_err(|error| classify_supervised_turn_failure(&error)); } () = &mut heartbeat => { *lease = self @@ -1624,7 +1636,7 @@ impl AiSupervisedAgentCoordinator { tokio::select! { result = &mut provider => { *lease = lease_state.lock().await.clone(); - return result.map_err(|_| SupervisedProviderTurnFailure::Provider); + return result.map_err(|error| classify_supervised_turn_failure(&error)); } () = &mut heartbeat => { let current = lease_state.lock().await.clone(); @@ -1700,9 +1712,22 @@ impl AiSupervisedAgentCoordinator { enum SupervisedProviderTurnFailure { Provider, + BudgetDenied, LeaseLost(AiError), } +/// Separates a certain pre-transport budget refusal from an uncertain turn. +/// +/// See the read-only coordinator for the full argument: the atomic budget +/// reservation happens before the transport boundary, so a denial proves no +/// bytes crossed it and no reservation was left held. +const fn classify_supervised_turn_failure(error: &AiError) -> SupervisedProviderTurnFailure { + match error { + AiError::PreTransportBudgetDenied => SupervisedProviderTurnFailure::BudgetDenied, + _ => SupervisedProviderTurnFailure::Provider, + } +} + #[cfg(test)] mod tests { use std::collections::VecDeque; @@ -2940,6 +2965,69 @@ mod tests { assert_eq!(run.final_states(), vec![AiRunState::Completed]); } + #[tokio::test] + async fn pre_transport_budget_denial_is_a_certain_supervised_failure() { + let lease = AiRunLease::test_running(principal_reference()); + let run = Arc::new(TestRunControl::new()); + let provider = Arc::new(TestProviderExecutor { + responses: Mutex::new(VecDeque::from([Err(AiError::PreTransportBudgetDenied)])), + require_checkpoint_cleared: false, + calls: AtomicUsize::new(0), + }); + let coordinator = AiSupervisedAgentCoordinator::new( + run.clone(), + provider.clone(), + Arc::new(TestOutputWriter), + Arc::new(TestCheckpointWriter { + provider_checkpoints: AtomicUsize::new(0), + }), + Arc::new(TestCheckpointControl { + adopted: Mutex::new(None), + consumed: AtomicBool::new(false), + }), + Arc::new(TestApprovalStager { + calls: AtomicUsize::new(0), + saw_checkpoint: AtomicBool::new(false), + }), + unused_automatic(), + unused_resume(), + Arc::new(TestRuleResolver), + Arc::new(TestPlanner { + scope: test_scope(), + route: test_route(), + continuation_count: AtomicUsize::new(0), + }), + Arc::new(FixedClock::new(time::OffsetDateTime::now_utc())), + limits(), + ); + + let outcome = coordinator + .execute_claimed(&lease) + .await + .expect("a supervised budget denial is a clean terminal failure"); + + assert_eq!( + outcome, + AiSupervisedAgentRunOutcome::Failed { + provider_turns: 0, + total_tool_calls: 0, + } + ); + assert_eq!(provider.calls.load(Ordering::SeqCst), 1); + assert_eq!(provider.remaining_responses(), 0); + assert_eq!(run.final_states(), vec![AiRunState::Failed]); + assert_eq!( + crate::classify_run_retry( + crate::AiRunRetryEvidence { + terminal: crate::AiRunTerminalEvent::Failed, + produced_assistant_output: false, + }, + Some("provider_budget_denied"), + ), + crate::AiRunRetryAdmission::Allowed + ); + } + #[tokio::test] async fn provider_turn_is_checkpointed_before_one_approval_is_staged() { let lease = AiRunLease::test_running(principal_reference()); diff --git a/crates/graphql-orm-ai/src/persistence.rs b/crates/graphql-orm-ai/src/persistence.rs index ca2b97a6..36eb9c58 100644 --- a/crates/graphql-orm-ai/src/persistence.rs +++ b/crates/graphql-orm-ai/src/persistence.rs @@ -454,7 +454,12 @@ pub(crate) struct AiBudgetCounterRecord { table = "graphql_orm_ai_budget_reservations", plural = "GraphqlOrmAiBudgetReservations", default_sort = "created_at DESC", - unique_composite = "principal_kind, principal_subject, idempotency_key" + unique_composite = "principal_kind, principal_subject, idempotency_key", + index( + name = "idx_graphql_orm_ai_budget_reservations_scope_state", + columns = ["scope_kind", "scope_id", "tenant_id", "state", "expires_at"], + directions = ["asc", "asc", "asc", "asc", "asc"] + ) )] #[cfg_attr(feature = "mssql", derive(GraphQLSchemaEntity))] #[cfg_attr( @@ -468,8 +473,11 @@ pub(crate) struct AiBudgetReservationRecord { pub id: graphql_orm::uuid::Uuid, #[graphql_orm(json, read = false, filter = false, order = false, subscribe = false)] pub budget_counter_ids: serde_json::Value, + #[filterable(type = "string")] pub scope_kind: String, + #[filterable(type = "string")] pub scope_id: String, + #[filterable(type = "string")] pub tenant_id: Option, #[filterable(type = "string")] pub principal_kind: String, @@ -502,6 +510,7 @@ pub(crate) struct AiBudgetReservationRecord { pub idempotency_key: String, #[filterable(type = "string")] pub state: String, + #[filterable(type = "number")] pub expires_at: i64, #[sortable] pub created_at: i64, @@ -2424,7 +2433,7 @@ pub(crate) struct AiRuntimeRecoveryRecord { /// Stable schema module ID. pub const AI_SCHEMA_MODULE_ID: &str = "com.dastari.graphql-orm-ai"; /// Current AI schema module version. -pub const AI_SCHEMA_MODULE_VERSION: &str = "0.62.0"; +pub const AI_SCHEMA_MODULE_VERSION: &str = "0.63.0"; /// Reserved table namespace. pub const AI_TABLE_NAMESPACE: &str = "graphql_orm_ai_"; diff --git a/crates/graphql-orm-ai/src/provider_calls.rs b/crates/graphql-orm-ai/src/provider_calls.rs index 18872c6a..b0fd9b3d 100644 --- a/crates/graphql-orm-ai/src/provider_calls.rs +++ b/crates/graphql-orm-ai/src/provider_calls.rs @@ -2898,6 +2898,7 @@ impl AiProviderCallExecutor { result.provider_session_claim = Some(current_claim); Ok(result) } + Err(error @ AiError::PreTransportBudgetDenied) => Err(error), Err(error) => { let _ = session_service .require_cleanup(¤t_claim, "provider_session_turn_ambiguous") @@ -2955,23 +2956,35 @@ impl AiProviderCallExecutor { return Err(AiError::Forbidden); } - let reservation = self + let reservation = match self .budget_service .reserve(&principal, plan.budget.clone()) - .await?; - let authorized_budget = reservation - .authorize_provider_call_with_reasoning_effort( - lease.run_id(), - lease.attempt_id(), - lease.lease_generation(), - &plan.provider_kind, - &plan.request.model, - plan.request.reasoning_effort, - plan.request.maximum_output_tokens.unwrap_or(0), - plan.request.maximum_builtin_tool_calls(), - self.clock.now(), - ) - .map_err(|_| AiError::BudgetDenied)?; + .await + { + Ok(reservation) => reservation, + Err(AiError::BudgetDenied) => return Err(AiError::PreTransportBudgetDenied), + Err(error) => return Err(error), + }; + let authorized_budget = match reservation.authorize_provider_call_with_reasoning_effort( + lease.run_id(), + lease.attempt_id(), + lease.lease_generation(), + &plan.provider_kind, + &plan.request.model, + plan.request.reasoning_effort, + plan.request.maximum_output_tokens.unwrap_or(0), + plan.request.maximum_builtin_tool_calls(), + self.clock.now(), + ) { + Ok(authorized) => authorized, + Err(_) => { + // The reservation exists but nothing was dispatched, so this is + // a provable release rather than capacity that would sit in the + // reserved column until its policy period rolled. + self.release_unstarted(&lease, &reservation).await?; + return Err(AiError::PreTransportBudgetDenied); + } + }; let mut context = match self .authorize_and_audit_transfers(&lease, &plan, authorized_budget) @@ -2980,7 +2993,10 @@ impl AiProviderCallExecutor { Ok(context) => context, Err(error) => { self.release_unstarted(&lease, &reservation).await?; - return Err(error); + return Err(match error { + AiError::BudgetDenied => AiError::PreTransportBudgetDenied, + error => error, + }); } }; @@ -3118,12 +3134,18 @@ impl AiProviderCallExecutor { .and_then(|calls| usize::try_from(calls).ok()) .unwrap_or(self.limits.maximum_builtin_tool_calls); let request_snapshot = plan.request.clone(); - let model_inference_manifest = plan + let model_inference_manifest = match plan .transfers .iter() .find(|manifest| manifest.capability == AiEgressCapability::ModelInference) .cloned() - .ok_or(AiError::EgressDenied)?; + { + Some(manifest) => manifest, + None => { + self.release_unstarted(&lease, &reservation).await?; + return Err(AiError::EgressDenied); + } + }; let previous_response_id = plan.request .continuation @@ -8687,6 +8709,98 @@ mod tests { .expect("reservation state should query") } + async fn reservation_count(database: &Database) -> i64 { + database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .count() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("reservation count should query") + } + + #[tokio::test] + async fn budget_denial_happens_before_transport_and_leaves_no_reservation() { + let fixture = fixture(Vec::new()).await; + let policy = fixture + .database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .limit(2) + .fetch_all() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("policy query should succeed") + .into_iter() + .next() + .expect("the fixture seeds one budget policy"); + AiBudgetPolicyRecord::update_by_id( + &fixture.database, + &policy.id, + UpdateAiBudgetPolicyRecordInput { + maximum_input_tokens: Some(Some(10)), + ..Default::default() + }, + ) + .await + .expect("policy ceiling should narrow below the planned estimate"); + + let executor = AiProviderCallExecutor::new( + fixture.runtime.clone(), + fixture.budget_service.clone(), + fixture.audit.clone(), + Arc::new(TestUsageAccounting), + Arc::new(SystemClock), + AiProviderCallLimits::new(64, 8_192, 64 * 1_024) + .expect("test provider limits should validate"), + ); + let error = executor + .execute(&fixture.lease, plan(&fixture)) + .await + .expect_err("an over-ceiling reservation must be denied"); + + assert!(matches!(error, AiError::PreTransportBudgetDenied)); + assert_eq!(fixture.mock.request_count(), 0); + assert_eq!(reservation_count(&fixture.database).await, 0); + } + + #[tokio::test] + async fn post_reservation_budget_binding_denial_releases_before_transport() { + let fixture = fixture(Vec::new()).await; + let mut plan = plan(&fixture); + // The plan was valid when authored, but this adversarial mutation asks + // the executor to authorize more output than the reservation holds. + // The executor must release the already-created reservation before it + // returns the proof-bearing pre-transport denial. + plan.request.maximum_output_tokens = Some(101); + + let executor = AiProviderCallExecutor::new( + fixture.runtime.clone(), + fixture.budget_service.clone(), + fixture.audit.clone(), + Arc::new(TestUsageAccounting), + Arc::new(SystemClock), + AiProviderCallLimits::new(64, 8_192, 64 * 1_024) + .expect("test provider limits should validate"), + ); + let error = executor + .execute(&fixture.lease, plan) + .await + .expect_err("a reservation that does not authorize the request must be denied"); + + assert!(matches!(error, AiError::PreTransportBudgetDenied)); + assert_eq!(fixture.mock.request_count(), 0); + assert_eq!(reservation_state(&fixture.database).await, "released"); + } + #[tokio::test] async fn successful_mock_turn_audits_egress_and_commits_authoritative_usage() { let fixture = fixture(vec![ @@ -9063,6 +9177,145 @@ mod tests { ); } + #[tokio::test] + async fn retained_pre_transport_budget_denial_does_not_require_cursor_cleanup() { + let cursor = AiProviderSessionCursor::new("mock.thread", "budget-denied-empty-thread") + .expect("test cursor should validate"); + let fixture = fixture_with_provider( + MockProvider::new(Vec::new()).with_provider_session_cursor(cursor), + ) + .await; + let session = AiSessionRecord::find_by_id(&fixture.database, &fixture.lease.session_id().0) + .await + .expect("session lookup should succeed") + .expect("session should exist"); + let update = AiSessionRecord::compare_and_swap( + &fixture.database, + &session.id, + session.row_version, + AiSessionRecordWhereInput::default(), + UpdateAiSessionRecordInput { + message_head: Some(1), + ..Default::default() + }, + ) + .await + .expect("session watermark update should succeed"); + assert!(matches!(update, ConditionalUpdateOutcome::Updated(_))); + AiMessageRecord::insert( + &fixture.database, + CreateAiMessageRecordInput { + id: fixture.lease.input_message_id(), + session_id: fixture.lease.session_id().0, + sequence: 1, + message_role: "user".to_owned(), + author_principal_kind: Some("user".to_owned()), + author_subject: Some(fixture.principal.subject().to_owned()), + client_message_id: Some(Uuid::new_v4()), + content_hash: Some("c".repeat(64)), + run_id: Some(fixture.lease.run_id().0), + provider_kind: None, + provider_model: None, + protected_preview: None, + block_count: 1, + completion_state: "complete".to_owned(), + finalized_at: Some(OffsetDateTime::now_utc().unix_timestamp()), + content_purged_at: None, + }, + ) + .await + .expect("input message should insert"); + let policy = fixture + .database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .limit(2) + .fetch_all() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("policy query should succeed") + .into_iter() + .next() + .expect("the fixture seeds one budget policy"); + AiBudgetPolicyRecord::update_by_id( + &fixture.database, + &policy.id, + UpdateAiBudgetPolicyRecordInput { + maximum_input_tokens: Some(Some(10)), + ..Default::default() + }, + ) + .await + .expect("policy ceiling should narrow below the planned estimate"); + + let provider_sessions = Arc::new( + OrmAiProviderSessionService::new( + fixture.database.clone(), + Arc::new(AllowAccess), + Arc::new(ProtectionPolicy), + Arc::new(DatabaseManagedContentProtector), + Arc::new(Resolver(fixture.principal.clone())), + Arc::new(SystemClock), + AiProviderSessionLimits::default(), + Duration::minutes(5), + ) + .expect("provider-session service should validate"), + ); + let descriptor = AiProviderSessionDescriptor::new( + ProviderKind::OpenAiCompatible, + "mock-profile", + "mock-model", + "a".repeat(64), + "mock-retained/v1", + "b".repeat(64), + ) + .expect("descriptor should validate"); + let executor = AiProviderCallExecutor::new( + fixture.runtime.clone(), + fixture.budget_service.clone(), + fixture.audit.clone(), + Arc::new(TestUsageAccounting), + Arc::new(SystemClock), + AiProviderCallLimits::new(64, 8_192, 64 * 1_024) + .expect("provider limits should validate"), + ); + + let error = executor + .execute_with_provider_session( + Arc::new(Mutex::new(fixture.lease.clone())), + plan(&fixture), + AiProviderSessionTurnPlan::new(descriptor, "d".repeat(64)) + .expect("session plan should validate"), + provider_sessions, + None, + ) + .await + .expect_err("the local budget denial must remain certain"); + + assert!(matches!(error, AiError::PreTransportBudgetDenied)); + assert_eq!(fixture.mock.request_count(), 0); + let bindings = fixture + .database + .transaction(TransactionMode::Default, |tx| { + Box::pin(async move { + tx.query::() + .limit(2) + .fetch_all() + .await + .map_err(OrmPublicError::from) + }) + }) + .await + .expect("provider-session binding query should succeed"); + assert_eq!(bindings.len(), 1); + assert_eq!(bindings[0].state, AiProviderSessionState::Claimed.as_str()); + assert!(bindings[0].cleanup_reason_code.is_none()); + } + #[tokio::test] async fn rebind_fence_loss_discards_the_exact_new_empty_provider_session() { let cursor = AiProviderSessionCursor::new("mock.thread", "losing-rebind-thread") diff --git a/crates/graphql-orm-ai/src/run_state.rs b/crates/graphql-orm-ai/src/run_state.rs index 5232efaf..d5602f4c 100644 --- a/crates/graphql-orm-ai/src/run_state.rs +++ b/crates/graphql-orm-ai/src/run_state.rs @@ -400,6 +400,7 @@ const fn is_retryable_failure_code(code: &str) -> bool { matches!( code.as_bytes(), b"provider_session_cleanup_unavailable" + | b"provider_budget_denied" | b"agent_rule_budget_exceeded" | b"agent_rule_changed_after_provider" | b"agent_turn_limit_reached" diff --git a/crates/graphql-orm-ai/tests/graphql_naming.rs b/crates/graphql-orm-ai/tests/graphql_naming.rs index ac7ba3ed..8e5bbca3 100644 --- a/crates/graphql-orm-ai/tests/graphql_naming.rs +++ b/crates/graphql-orm-ai/tests/graphql_naming.rs @@ -38,6 +38,10 @@ fn configured_graphql_case_is_coherent_without_aliases() { assert!(configuration_sdl.contains("LOCAL_HARNESS")); assert!(configuration_sdl.contains("openaiCompatible: AiOpenAiCompatibleProfileInput")); assert!(configuration_sdl.contains("providerRetainedContinuation: Boolean!")); + assert!(configuration_sdl.contains("aiBudgetScopeCapacity(scope:")); + assert!(configuration_sdl.contains("reclaimAiBudgetReservation(input:")); + assert!(configuration_sdl.contains("uncertainReservationCount: Int!")); + assert!(configuration_sdl.contains("reclaimable: Boolean!")); assert!(skill_sdl.contains("aiSkills(scope:")); assert!(skill_sdl.contains("upsertAiSkill(input:")); assert!(skill_sdl.contains("publishAiSkillVersion(input:")); @@ -64,6 +68,10 @@ fn configured_graphql_case_is_coherent_without_aliases() { assert!(configuration_sdl.contains("LocalHarness")); assert!(configuration_sdl.contains("OpenaiCompatible: AiOpenAiCompatibleProfileInput")); assert!(configuration_sdl.contains("ProviderRetainedContinuation: Boolean!")); + assert!(configuration_sdl.contains("AiBudgetScopeCapacity(Scope:")); + assert!(configuration_sdl.contains("ReclaimAiBudgetReservation(Input:")); + assert!(configuration_sdl.contains("UncertainReservationCount: Int!")); + assert!(configuration_sdl.contains("Reclaimable: Boolean!")); assert!(skill_sdl.contains("AiSkills(Scope:")); assert!(skill_sdl.contains("UpsertAiSkill(Input:")); assert!(skill_sdl.contains("PublishAiSkillVersion(Input:")); diff --git a/crates/graphql-orm-ai/tests/schema_module.rs b/crates/graphql-orm-ai/tests/schema_module.rs index 5ae69679..d6b40f4b 100644 --- a/crates/graphql-orm-ai/tests/schema_module.rs +++ b/crates/graphql-orm-ai/tests/schema_module.rs @@ -8,7 +8,7 @@ fn ai_schema_module_owns_only_reserved_namespace_tables() { assert_eq!(catalog.modules().len(), 1); assert_eq!(catalog.modules()[0].version, AI_SCHEMA_MODULE_VERSION); - assert_eq!(AI_SCHEMA_MODULE_VERSION, "0.62.0"); + assert_eq!(AI_SCHEMA_MODULE_VERSION, "0.63.0"); assert_eq!(catalog.entities().len(), 47); assert!( catalog @@ -107,6 +107,11 @@ fn ai_schema_module_owns_only_reserved_namespace_tables() { "idempotency_key".to_owned(), ] })); + assert!(reservation.indexes.iter().any(|index| { + index.name == "idx_graphql_orm_ai_budget_reservations_scope_state" + && index.columns == ["scope_kind", "scope_id", "tenant_id", "state", "expires_at"] + && !index.is_unique + })); assert!( reservation .columns diff --git a/docs/reference/workspace-packages.md b/docs/reference/workspace-packages.md index bb93e2e9..be9ff6d9 100644 --- a/docs/reference/workspace-packages.md +++ b/docs/reference/workspace-packages.md @@ -19,8 +19,8 @@ changes. | Package | Version | Path | Default features | Direct internal dependencies | | --- | --- | --- | --- | --- | | `graphql-orm` | `0.23.0` | `crates/graphql-orm` | `sqlite` | `graphql-orm-macros`, `graphql-orm-operation-catalog`, `graphql-orm-router-protocol` (dev-only) | -| `graphql-orm-ai` | `0.83.0` | `crates/graphql-orm-ai` | `sqlite` | `graphql-orm`, `graphql-orm-ai-tool-profiles`, `graphql-orm-storage` | -| `graphql-orm-ai-tool-profiles` | `0.6.0` | `crates/graphql-orm-ai-tool-profiles` | none | `graphql-orm-operation-catalog`, `graphql-orm-router-protocol` (dev-only) | +| `graphql-orm-ai` | `0.84.0` | `crates/graphql-orm-ai` | `sqlite` | `graphql-orm`, `graphql-orm-ai-tool-profiles`, `graphql-orm-storage` | +| `graphql-orm-ai-tool-profiles` | `0.7.0` | `crates/graphql-orm-ai-tool-profiles` | none | `graphql-orm-operation-catalog`, `graphql-orm-router-protocol` (dev-only) | | `graphql-orm-backup` | `0.7.1` | `crates/graphql-orm-backup` | `local` | `graphql-orm` (optional), `graphql-orm-storage` | | `graphql-orm-macros` | `0.23.0` | `crates/graphql-orm-macros` | `sqlite` | none | | `graphql-orm-operation-catalog` | `0.3.0` | `crates/graphql-orm-operation-catalog` | none | `graphql-orm-router-protocol` (optional) |