codexLB-up 1.25.0b10__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (827) hide show
  1. app/__init__.py +10 -0
  2. app/admin_cli.py +494 -0
  3. app/cli.py +339 -0
  4. app/codex_sessions_retag.py +487 -0
  5. app/core/__init__.py +0 -0
  6. app/core/audit/__init__.py +0 -0
  7. app/core/audit/service.py +216 -0
  8. app/core/audit/types.py +94 -0
  9. app/core/auth/__init__.py +236 -0
  10. app/core/auth/api_key_cache.py +65 -0
  11. app/core/auth/dashboard_access.py +429 -0
  12. app/core/auth/dashboard_mode.py +171 -0
  13. app/core/auth/dashboard_session_ttl.py +74 -0
  14. app/core/auth/dashboard_users_cache.py +120 -0
  15. app/core/auth/dependencies.py +710 -0
  16. app/core/auth/external_identity.py +78 -0
  17. app/core/auth/guardian.py +374 -0
  18. app/core/auth/models.py +49 -0
  19. app/core/auth/providers/__init__.py +129 -0
  20. app/core/auth/providers/oidc.py +838 -0
  21. app/core/auth/providers/registry.py +118 -0
  22. app/core/auth/refresh.py +394 -0
  23. app/core/auth/step_up.py +117 -0
  24. app/core/auth/totp.py +81 -0
  25. app/core/balancer/__init__.py +71 -0
  26. app/core/balancer/logic.py +1612 -0
  27. app/core/balancer/rendezvous_hash.py +21 -0
  28. app/core/balancer/types.py +21 -0
  29. app/core/bootstrap.py +172 -0
  30. app/core/cache/invalidation.py +426 -0
  31. app/core/clients/__init__.py +0 -0
  32. app/core/clients/codex.py +996 -0
  33. app/core/clients/codex_version.py +187 -0
  34. app/core/clients/codex_version_snapshot.py +2 -0
  35. app/core/clients/files.py +534 -0
  36. app/core/clients/headers.py +29 -0
  37. app/core/clients/http.py +511 -0
  38. app/core/clients/model_fetcher.py +276 -0
  39. app/core/clients/native_egress.py +1304 -0
  40. app/core/clients/oauth.py +477 -0
  41. app/core/clients/proxy.py +6089 -0
  42. app/core/clients/proxy_websocket.py +1411 -0
  43. app/core/clients/rate_limit_reset_credits.py +412 -0
  44. app/core/clients/stream_errors.py +12 -0
  45. app/core/clients/thread_cache_identity.py +347 -0
  46. app/core/clients/usage.py +467 -0
  47. app/core/clock.py +170 -0
  48. app/core/config/__init__.py +0 -0
  49. app/core/config/background_jobs.py +73 -0
  50. app/core/config/context_window_overrides.py +166 -0
  51. app/core/config/dashboard_overrides.py +145 -0
  52. app/core/config/inheritable.py +47 -0
  53. app/core/config/key_fingerprint.py +115 -0
  54. app/core/config/settings.py +872 -0
  55. app/core/config/settings_cache.py +100 -0
  56. app/core/config/spool_retention.py +217 -0
  57. app/core/config/tiers.py +174 -0
  58. app/core/conversation_archive.py +568 -0
  59. app/core/crypto.py +37 -0
  60. app/core/draining.py +11 -0
  61. app/core/errors.py +365 -0
  62. app/core/exceptions.py +155 -0
  63. app/core/handlers/__init__.py +3 -0
  64. app/core/handlers/exceptions.py +478 -0
  65. app/core/http_protocol.py +176 -0
  66. app/core/http_protocol_httptools.py +204 -0
  67. app/core/ingress_limits.py +17 -0
  68. app/core/metadata_store.py +18 -0
  69. app/core/metrics/__init__.py +0 -0
  70. app/core/metrics/middleware.py +74 -0
  71. app/core/metrics/prometheus.py +604 -0
  72. app/core/middleware/__init__.py +25 -0
  73. app/core/middleware/api_firewall.py +111 -0
  74. app/core/middleware/app_version.py +37 -0
  75. app/core/middleware/dashboard_auth_proxy.py +55 -0
  76. app/core/middleware/dashboard_csrf.py +177 -0
  77. app/core/middleware/dashboard_gzip.py +41 -0
  78. app/core/middleware/dashboard_overrides.py +59 -0
  79. app/core/middleware/firewall_cache.py +91 -0
  80. app/core/middleware/inflight.py +126 -0
  81. app/core/middleware/multipart_content_encoding.py +129 -0
  82. app/core/middleware/path_rewrite.py +148 -0
  83. app/core/middleware/request_body_limit.py +196 -0
  84. app/core/middleware/request_decompression.py +209 -0
  85. app/core/middleware/request_id.py +72 -0
  86. app/core/middleware/required_capability_http.py +116 -0
  87. app/core/middleware/trusted_proxy_headers.py +36 -0
  88. app/core/multipart.py +237 -0
  89. app/core/multipart_fields.py +91 -0
  90. app/core/openai/__init__.py +0 -0
  91. app/core/openai/chat_requests.py +493 -0
  92. app/core/openai/chat_responses.py +931 -0
  93. app/core/openai/contracts.py +46 -0
  94. app/core/openai/exceptions.py +24 -0
  95. app/core/openai/host_models.py +14 -0
  96. app/core/openai/images.py +413 -0
  97. app/core/openai/message_coercion.py +376 -0
  98. app/core/openai/model_refresh_scheduler.py +423 -0
  99. app/core/openai/model_registry.py +997 -0
  100. app/core/openai/model_registry_store.py +548 -0
  101. app/core/openai/models.py +172 -0
  102. app/core/openai/parsing.py +97 -0
  103. app/core/openai/requests.py +1980 -0
  104. app/core/openai/strict_schema.py +199 -0
  105. app/core/openai/tool_call_safety.py +99 -0
  106. app/core/openai/v1_requests.py +187 -0
  107. app/core/plan_types.py +90 -0
  108. app/core/prestop.py +202 -0
  109. app/core/rate_limiter/__init__.py +0 -0
  110. app/core/rate_limiter/db_rate_limiter.py +215 -0
  111. app/core/request_locality.py +473 -0
  112. app/core/resilience/__init__.py +0 -0
  113. app/core/resilience/backpressure.py +57 -0
  114. app/core/resilience/bulkhead.py +202 -0
  115. app/core/resilience/circuit_breaker.py +150 -0
  116. app/core/resilience/degradation.py +55 -0
  117. app/core/resilience/loop_lag_monitor.py +60 -0
  118. app/core/resilience/memory_monitor.py +137 -0
  119. app/core/resilience/network_recovery.py +256 -0
  120. app/core/resilience/overload.py +98 -0
  121. app/core/resilience/toggles.py +114 -0
  122. app/core/retention/__init__.py +0 -0
  123. app/core/retention/job.py +240 -0
  124. app/core/retention/scheduler.py +69 -0
  125. app/core/runtime_logging.py +601 -0
  126. app/core/scheduling/__init__.py +0 -0
  127. app/core/scheduling/leader_election.py +993 -0
  128. app/core/scheduling/leader_election_handle.py +25 -0
  129. app/core/server.py +140 -0
  130. app/core/shutdown.py +326 -0
  131. app/core/socket_peer.py +22 -0
  132. app/core/startup.py +41 -0
  133. app/core/timeout_invariants.py +347 -0
  134. app/core/tracing/__init__.py +0 -0
  135. app/core/tracing/otel.py +207 -0
  136. app/core/types.py +6 -0
  137. app/core/upstream_proxy/__init__.py +13 -0
  138. app/core/upstream_proxy/cache.py +130 -0
  139. app/core/upstream_proxy/resolver.py +189 -0
  140. app/core/upstream_proxy/types.py +104 -0
  141. app/core/usage/__init__.py +407 -0
  142. app/core/usage/depletion.py +109 -0
  143. app/core/usage/live_hub.py +41 -0
  144. app/core/usage/live_snapshots.py +174 -0
  145. app/core/usage/logs.py +212 -0
  146. app/core/usage/metadata_scheduler.py +118 -0
  147. app/core/usage/models.py +54 -0
  148. app/core/usage/pricing.py +576 -0
  149. app/core/usage/pricing_catalog.py +241 -0
  150. app/core/usage/pricing_snapshot.json +821 -0
  151. app/core/usage/quota.py +133 -0
  152. app/core/usage/refresh_policy.py +21 -0
  153. app/core/usage/refresh_scheduler.py +713 -0
  154. app/core/usage/reset_credits_refresh_scheduler.py +435 -0
  155. app/core/usage/types.py +168 -0
  156. app/core/utils/__init__.py +0 -0
  157. app/core/utils/json_guards.py +18 -0
  158. app/core/utils/locks.py +33 -0
  159. app/core/utils/masking.py +20 -0
  160. app/core/utils/proxy_env.py +104 -0
  161. app/core/utils/request_id.py +65 -0
  162. app/core/utils/retry.py +77 -0
  163. app/core/utils/shared_future.py +172 -0
  164. app/core/utils/sse.py +286 -0
  165. app/core/utils/time.py +24 -0
  166. app/db/__init__.py +0 -0
  167. app/db/account_identity_lock.py +56 -0
  168. app/db/alembic/__init__.py +1 -0
  169. app/db/alembic/env.py +70 -0
  170. app/db/alembic/revision_ids.py +43 -0
  171. app/db/alembic/versions/20260213_000000_base_schema.py +248 -0
  172. app/db/alembic/versions/20260213_000100_normalize_account_plan_types.py +45 -0
  173. app/db/alembic/versions/20260213_000200_add_request_logs_reasoning_effort.py +45 -0
  174. app/db/alembic/versions/20260213_000300_add_accounts_reset_at.py +45 -0
  175. app/db/alembic/versions/20260213_000400_add_accounts_chatgpt_account_id.py +45 -0
  176. app/db/alembic/versions/20260213_000500_add_dashboard_settings.py +70 -0
  177. app/db/alembic/versions/20260213_000600_add_dashboard_settings_totp.py +64 -0
  178. app/db/alembic/versions/20260213_000700_add_dashboard_settings_password.py +59 -0
  179. app/db/alembic/versions/20260213_000800_add_api_keys.py +87 -0
  180. app/db/alembic/versions/20260214_000000_add_api_key_limits.py +143 -0
  181. app/db/alembic/versions/20260215_000000_add_idx_logs_requested_at.py +38 -0
  182. app/db/alembic/versions/20260218_000000_add_api_key_usage_reservations.py +117 -0
  183. app/db/alembic/versions/20260218_000100_add_import_without_overwrite_and_drop_accounts_email_unique.py +198 -0
  184. app/db/alembic/versions/20260225_000000_add_dashboard_settings_routing_strategy.py +53 -0
  185. app/db/alembic/versions/20260228_020000_align_api_key_limit_enum_types.py +104 -0
  186. app/db/alembic/versions/20260228_030000_add_api_firewall_allowlist.py +42 -0
  187. app/db/alembic/versions/20260306_000000_add_request_logs_service_tier.py +45 -0
  188. app/db/alembic/versions/20260307_000000_add_api_key_enforcement_fields.py +56 -0
  189. app/db/alembic/versions/20260308_000000_add_sqlite_performance_indexes.py +41 -0
  190. app/db/alembic/versions/20260309_000000_add_additional_usage_history.py +71 -0
  191. app/db/alembic/versions/20260309_000000_request_logs_nullable_account_id.py +28 -0
  192. app/db/alembic/versions/20260310_000000_add_request_logs_transport.py +45 -0
  193. app/db/alembic/versions/20260310_000000_fix_postgresql_enum_value_casing.py +79 -0
  194. app/db/alembic/versions/20260310_120000_add_sticky_session_kinds_and_affinity_ttl.py +110 -0
  195. app/db/alembic/versions/20260312_000000_add_additional_usage_quota_key.py +193 -0
  196. app/db/alembic/versions/20260312_000000_split_sticky_sessions_primary_key_by_kind.py +96 -0
  197. app/db/alembic/versions/20260312_010000_merge_additional_usage_and_sticky_session_heads.py +25 -0
  198. app/db/alembic/versions/20260312_020000_merge_request_logs_transport_and_feature_heads.py +25 -0
  199. app/db/alembic/versions/20260312_020000_merge_request_logs_transport_and_sticky_heads.py +25 -0
  200. app/db/alembic/versions/20260312_120000_add_dashboard_upstream_stream_transport.py +62 -0
  201. app/db/alembic/versions/20260313_024500_merge_request_log_transport_heads.py +27 -0
  202. app/db/alembic/versions/20260319_100937_increase_prompt_cache_ttl_to_1800s.py +73 -0
  203. app/db/alembic/versions/20260319_111500_merge_prompt_cache_ttl_and_upstream_transport_heads.py +25 -0
  204. app/db/alembic/versions/20260319_183000_normalize_sqlite_account_status_casing.py +48 -0
  205. app/db/alembic/versions/20260320_000000_add_request_log_requested_actual_tiers.py +51 -0
  206. app/db/alembic/versions/20260321_190000_tighten_dashboard_db_indexes.py +64 -0
  207. app/db/alembic/versions/20260321_210000_merge_request_log_tiers_and_dashboard_index_heads.py +25 -0
  208. app/db/alembic/versions/20260325_000000_add_request_log_cost.py +193 -0
  209. app/db/alembic/versions/20260328_000000_add_rate_limit_attempts.py +48 -0
  210. app/db/alembic/versions/20260328_120000_add_request_log_ttft.py +37 -0
  211. app/db/alembic/versions/20260328_130000_add_audit_logs_table.py +45 -0
  212. app/db/alembic/versions/20260328_140000_add_scheduler_leader_table.py +39 -0
  213. app/db/alembic/versions/20260330_000000_add_cache_locality_settings.py +69 -0
  214. app/db/alembic/versions/20260330_010000_merge_scheduler_leader_and_cache_locality_heads.py +25 -0
  215. app/db/alembic/versions/20260330_020000_add_bridge_ring_members.py +46 -0
  216. app/db/alembic/versions/20260401_000000_add_api_key_enforced_service_tier.py +52 -0
  217. app/db/alembic/versions/20260401_000000_add_cache_invalidation.py +41 -0
  218. app/db/alembic/versions/20260401_020000_merge_cache_invalidation_and_api_key_service_tier_heads.py +22 -0
  219. app/db/alembic/versions/20260402_000000_switch_dashboard_routing_default_to_capacity_weighted.py +66 -0
  220. app/db/alembic/versions/20260403_000000_add_credit_api_key_limit_values.py +44 -0
  221. app/db/alembic/versions/20260406_000000_add_api_key_account_assignments.py +57 -0
  222. app/db/alembic/versions/20260406_010000_add_api_key_assignment_scope_flag.py +75 -0
  223. app/db/alembic/versions/20260407_000000_add_bridge_gateway_safe_mode.py +51 -0
  224. app/db/alembic/versions/20260407_010000_merge_api_key_assignment_and_bridge_gateway_heads.py +25 -0
  225. app/db/alembic/versions/20260408_000000_switch_import_without_overwrite_default_to_true.py +73 -0
  226. app/db/alembic/versions/20260408_010000_merge_import_without_overwrite_and_assignment_heads.py +25 -0
  227. app/db/alembic/versions/20260409_000000_add_http_bridge_sessions.py +130 -0
  228. app/db/alembic/versions/20260409_000000_switch_sticky_threads_and_prefer_earlier_reset_defaults_to_true.py +89 -0
  229. app/db/alembic/versions/20260409_010000_add_dashboard_settings_bootstrap_token.py +57 -0
  230. app/db/alembic/versions/20260409_010000_merge_http_bridge_and_import_without_overwrite_heads.py +25 -0
  231. app/db/alembic/versions/20260409_020000_fix_http_bridge_last_seen_index.py +46 -0
  232. app/db/alembic/versions/20260410_000000_backfill_pristine_dashboard_settings_defaults.py +77 -0
  233. app/db/alembic/versions/20260410_030000_restore_import_without_overwrite_default_true.py +73 -0
  234. app/db/alembic/versions/20260410_040000_merge_dashboard_defaults_and_import_default_heads.py +25 -0
  235. app/db/alembic/versions/20260413_000000_add_accounts_blocked_at.py +44 -0
  236. app/db/alembic/versions/20260415_160000_add_automations_tables.py +189 -0
  237. app/db/alembic/versions/20260415_160000_add_request_logs_response_lookup_index.py +67 -0
  238. app/db/alembic/versions/20260417_000000_add_request_log_plan_type.py +44 -0
  239. app/db/alembic/versions/20260418_190000_add_automation_runs_cycle_key.py +87 -0
  240. app/db/alembic/versions/20260419_000000_add_automation_run_cycle_metadata.py +56 -0
  241. app/db/alembic/versions/20260419_020000_add_automation_run_cycles_snapshot_tables.py +342 -0
  242. app/db/alembic/versions/20260421_000000_add_dashboard_session_ttl_setting.py +51 -0
  243. app/db/alembic/versions/20260421_120000_merge_request_log_lookup_and_plan_type_heads.py +25 -0
  244. app/db/alembic/versions/20260421_130000_merge_automation_and_request_log_heads.py +25 -0
  245. app/db/alembic/versions/20260422_103000_normalize_legacy_automation_run_cycle_keys.py +648 -0
  246. app/db/alembic/versions/20260423_120000_add_api_key_limit_reset_at_index.py +43 -0
  247. app/db/alembic/versions/20260424_000000_merge_dashboard_session_ttl_and_request_log_heads.py +25 -0
  248. app/db/alembic/versions/20260426_000000_add_dashboard_relative_availability_settings.py +68 -0
  249. app/db/alembic/versions/20260427_130000_add_warmup_model_and_request_kind.py +75 -0
  250. app/db/alembic/versions/20260506_000000_add_account_routing_policy.py +37 -0
  251. app/db/alembic/versions/20260509_000000_add_api_key_traffic_class.py +37 -0
  252. app/db/alembic/versions/20260509_010000_add_additional_quota_routing_policies.py +37 -0
  253. app/db/alembic/versions/20260513_000000_add_accounts_alias.py +44 -0
  254. app/db/alembic/versions/20260513_000000_add_api_key_apply_to_codex_model.py +59 -0
  255. app/db/alembic/versions/20260514_000000_add_request_logs_api_key_time_index.py +43 -0
  256. app/db/alembic/versions/20260515_000000_add_prefer_earlier_reset_window.py +47 -0
  257. app/db/alembic/versions/20260515_000000_soft_delete_request_logs_on_account_delete.py +83 -0
  258. app/db/alembic/versions/20260515_020000_add_split_sticky_budget_thresholds.py +62 -0
  259. app/db/alembic/versions/20260516_000000_add_sqlite_hot_path_indexes.py +75 -0
  260. app/db/alembic/versions/20260517_000000_merge_sqlite_hot_path_and_request_log_delete_heads.py +24 -0
  261. app/db/alembic/versions/20260518_000000_add_http_bridge_durable_input_prefix.py +44 -0
  262. app/db/alembic/versions/20260518_010000_merge_http_bridge_and_sqlite_recovery_heads.py +25 -0
  263. app/db/alembic/versions/20260518_120000_add_dashboard_guest_access.py +57 -0
  264. app/db/alembic/versions/20260520_000000_merge_api_key_and_http_bridge_heads.py +26 -0
  265. app/db/alembic/versions/20260520_010000_add_request_logs_api_key_account_index.py +45 -0
  266. app/db/alembic/versions/20260520_010000_merge_automation_and_api_key_heads.py +26 -0
  267. app/db/alembic/versions/20260520_030000_add_quota_planner.py +194 -0
  268. app/db/alembic/versions/20260520_040000_merge_quota_planner_and_request_log_index_heads.py +24 -0
  269. app/db/alembic/versions/20260521_000000_add_account_security_work_authorized.py +46 -0
  270. app/db/alembic/versions/20260521_000000_merge_dashboard_guest_and_api_key_heads.py +26 -0
  271. app/db/alembic/versions/20260521_000000_merge_prefer_reset_window_head.py +26 -0
  272. app/db/alembic/versions/20260521_000000_merge_warmup_and_request_log_index_heads.py +25 -0
  273. app/db/alembic/versions/20260522_000000_add_limit_warmup_trigger.py +195 -0
  274. app/db/alembic/versions/20260524_000000_merge_accounts_alias_and_reset_window_heads.py +24 -0
  275. app/db/alembic/versions/20260525_000000_add_usage_raw_window_latest_index.py +31 -0
  276. app/db/alembic/versions/20260526_000000_add_request_log_failure_metadata.py +100 -0
  277. app/db/alembic/versions/20260531_000000_add_account_workspace_identity.py +39 -0
  278. app/db/alembic/versions/20260531_000000_add_single_account_routing.py +42 -0
  279. app/db/alembic/versions/20260601_000000_merge_relative_availability_and_usage_raw_heads.py +24 -0
  280. app/db/alembic/versions/20260601_010000_add_api_key_usage_sections.py +58 -0
  281. app/db/alembic/versions/20260601_010000_add_routing_policy_persistence.py +67 -0
  282. app/db/alembic/versions/20260601_010000_merge_security_work_and_relative_availability_heads.py +25 -0
  283. app/db/alembic/versions/20260601_010000_merge_split_sticky_and_relative_usage_heads.py +25 -0
  284. app/db/alembic/versions/20260601_010000_merge_warmup_and_relative_availability_heads.py +25 -0
  285. app/db/alembic/versions/20260601_020000_merge_additional_quota_routing_and_relative_availability_heads.py +25 -0
  286. app/db/alembic/versions/20260601_020000_merge_quota_planner_and_relative_availability_heads.py +25 -0
  287. app/db/alembic/versions/20260601_020000_merge_reset_window_and_relative_availability_heads.py +25 -0
  288. app/db/alembic/versions/20260601_020000_merge_routing_policy_and_split_threshold_heads.py +24 -0
  289. app/db/alembic/versions/20260601_020000_merge_warmup_and_request_log_failure_heads.py +25 -0
  290. app/db/alembic/versions/20260601_030000_merge_single_account_and_routing_quota_heads.py +29 -0
  291. app/db/alembic/versions/20260601_040000_merge_security_work_and_routing_consolidation_heads.py +25 -0
  292. app/db/alembic/versions/20260601_050000_merge_quota_planner_and_routing_consolidation_heads.py +25 -0
  293. app/db/alembic/versions/20260602_050000_add_upstream_proxy_routing.py +164 -0
  294. app/db/alembic/versions/20260602_060000_merge_account_workspace_and_failure_heads.py +25 -0
  295. app/db/alembic/versions/20260602_070000_add_request_log_useragent_fields.py +86 -0
  296. app/db/alembic/versions/20260602_070000_merge_quota_planner_and_account_workspace_heads.py +25 -0
  297. app/db/alembic/versions/20260602_080000_merge_upstream_proxy_and_quota_planner_heads.py +25 -0
  298. app/db/alembic/versions/20260602_080000_merge_useragent_and_upstream_proxy_heads.py +24 -0
  299. app/db/alembic/versions/20260603_000000_add_weekly_pace_working_days.py +52 -0
  300. app/db/alembic/versions/20260603_000000_free_account_monthly_window.py +60 -0
  301. app/db/alembic/versions/20260604_000000_add_reauth_required_account_status.py +159 -0
  302. app/db/alembic/versions/20260604_010000_merge_automations_and_reauth_status_heads.py +25 -0
  303. app/db/alembic/versions/20260604_010000_merge_dashboard_guest_and_reauth_status_heads.py +25 -0
  304. app/db/alembic/versions/20260604_010000_merge_useragent_and_reauth_heads.py +24 -0
  305. app/db/alembic/versions/20260607_000000_merge_weekly_monthly_useragent_heads.py +28 -0
  306. app/db/alembic/versions/20260608_000000_add_hide_upstream_quota_from_api_keys.py +50 -0
  307. app/db/alembic/versions/20260611_000000_add_request_log_archive_request_id.py +47 -0
  308. app/db/alembic/versions/20260611_000000_merge_automations_and_weekly_monthly_useragent_heads.py +26 -0
  309. app/db/alembic/versions/20260611_000000_merge_dashboard_guest_and_weekly_useragent_heads.py +26 -0
  310. app/db/alembic/versions/20260611_010000_add_request_log_client_ip.py +57 -0
  311. app/db/alembic/versions/20260611_020000_merge_request_log_archive_and_client_ip_heads.py +26 -0
  312. app/db/alembic/versions/20260611_030000_merge_dashboard_guest_and_request_log_client_ip_heads.py +26 -0
  313. app/db/alembic/versions/20260613_000000_add_accounts_codex_installation_id.py +57 -0
  314. app/db/alembic/versions/20260613_010000_merge_installation_id_and_dashboard_weekly_heads.py +26 -0
  315. app/db/alembic/versions/20260615_000000_merge_automations_and_dashboard_guest_heads.py +26 -0
  316. app/db/alembic/versions/20260623_000000_add_limit_warmup_exhausted_threshold.py +51 -0
  317. app/db/alembic/versions/20260626_000000_add_http_downstream_transport_policy.py +46 -0
  318. app/db/alembic/versions/20260626_010000_add_request_logs_upstream_transport.py +31 -0
  319. app/db/alembic/versions/20260629_000000_add_dashboard_query_hot_path_indexes.py +53 -0
  320. app/db/alembic/versions/20260630_000000_add_limit_warmup_staggered_idle.py +47 -0
  321. app/db/alembic/versions/20260630_000000_merge_request_log_client_ip_and_dashboard_hot_path_heads.py +26 -0
  322. app/db/alembic/versions/20260630_010000_merge_warmup_and_request_log_dashboard_heads.py +26 -0
  323. app/db/alembic/versions/20260630_020000_merge_automations_and_main_heads.py +26 -0
  324. app/db/alembic/versions/20260630_020000_merge_warmup_threshold_and_main_heads.py +26 -0
  325. app/db/alembic/versions/20260630_030000_add_automation_run_model_snapshot.py +59 -0
  326. app/db/alembic/versions/20260630_040000_merge_automation_snapshot_and_warmup_threshold_heads.py +26 -0
  327. app/db/alembic/versions/20260630_050000_add_automation_run_prompt_snapshot.py +52 -0
  328. app/db/alembic/versions/20260701_000000_add_weekly_pace_smoothing_minutes.py +51 -0
  329. app/db/alembic/versions/20260702_000000_add_model_sources.py +176 -0
  330. app/db/alembic/versions/20260705_000000_add_model_source_audio_transcriptions.py +50 -0
  331. app/db/alembic/versions/20260705_000000_harden_dashboard_session_ttl.py +62 -0
  332. app/db/alembic/versions/20260706_000000_add_model_source_audio_per_minute.py +43 -0
  333. app/db/alembic/versions/20260707_000000_merge_model_sources_and_session_ttl.py +24 -0
  334. app/db/alembic/versions/20260707_010000_merge_automations_and_model_sources_heads.py +24 -0
  335. app/db/alembic/versions/20260707_020000_add_automation_job_account_scope.py +125 -0
  336. app/db/alembic/versions/20260709_000000_add_ttft_phase_observability.py +66 -0
  337. app/db/alembic/versions/20260711_000000_add_dashboard_prohibit_fast_mode.py +42 -0
  338. app/db/alembic/versions/20260711_010000_add_account_chatgpt_user_id.py +40 -0
  339. app/db/alembic/versions/20260711_020000_add_dashboard_account_concurrency_caps.py +71 -0
  340. app/db/alembic/versions/20260711_030000_add_limit_warmup_idle_threshold.py +51 -0
  341. app/db/alembic/versions/20260712_010000_add_account_usage_rollups.py +56 -0
  342. app/db/alembic/versions/20260712_020000_add_api_key_usage_rollups.py +53 -0
  343. app/db/alembic/versions/20260713_020000_add_model_registry_snapshot.py +42 -0
  344. app/db/alembic/versions/20260713_040000_add_account_refresh_claims.py +42 -0
  345. app/db/alembic/versions/20260713_040000_add_replica_guardrails.py +64 -0
  346. app/db/alembic/versions/20260713_070000_add_reset_credit_redeem_tables.py +58 -0
  347. app/db/alembic/versions/20260714_000000_add_oauth_flow_states.py +63 -0
  348. app/db/alembic/versions/20260715_000000_add_request_log_queue_latency.py +44 -0
  349. app/db/alembic/versions/20260716_000000_add_oauth_device_flow_slots.py +41 -0
  350. app/db/alembic/versions/20260716_010000_add_dashboard_retention_settings.py +48 -0
  351. app/db/alembic/versions/20260716_010000_add_reset_credit_display_settings.py +73 -0
  352. app/db/alembic/versions/20260717_000000_add_http_bridge_retry_circuits.py +47 -0
  353. app/db/alembic/versions/20260717_000000_merge_retention_and_reset_credit_display_heads.py +26 -0
  354. app/db/alembic/versions/20260717_000000_optimize_dashboard_hot_path_indexes.py +162 -0
  355. app/db/alembic/versions/20260717_000001_merge_retry_circuits_and_dashboard_indexes.py +22 -0
  356. app/db/alembic/versions/20260720_000000_add_request_log_conversation_id.py +57 -0
  357. app/db/alembic/versions/20260722_000000_backfill_request_log_useragent_families.py +40 -0
  358. app/db/alembic/versions/20260724_000000_add_request_usage_time_rollups.py +176 -0
  359. app/db/alembic/versions/20260724_000000_merge_request_log_schema_heads.py +33 -0
  360. app/db/alembic/versions/20260725_000000_add_http_bridge_pending_tool_calls.py +43 -0
  361. app/db/alembic/versions/20260726_000000_add_account_plan_downgrade_observations.py +44 -0
  362. app/db/alembic/versions/20260726_000000_repair_request_usage_rollups_after_merge.py +37 -0
  363. app/db/alembic/versions/20260727_000000_add_sticky_session_continuity_abandoned_at.py +45 -0
  364. app/db/alembic/versions/20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads.py +31 -0
  365. app/db/alembic/versions/20260730_000000_add_api_key_fair_share_threshold.py +44 -0
  366. app/db/alembic/versions/20260730_000000_add_http_bridge_recovery_attempts.py +60 -0
  367. app/db/alembic/versions/20260731_000000_add_capability_lineage_markers.py +38 -0
  368. app/db/alembic/versions/20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads.py +30 -0
  369. app/db/alembic/versions/20260804_000000_add_http_bridge_operations.py +63 -0
  370. app/db/alembic/versions/20260804_000001_add_global_http_bridge_operation_fingerprint.py +53 -0
  371. app/db/alembic/versions/20260804_230000_add_request_log_connection_request_kind.py +43 -0
  372. app/db/alembic/versions/20260805_000000_add_http_bridge_operation_spool.py +75 -0
  373. app/db/alembic/versions/20260805_000001_finalize_http_bridge_operation_spool.py +65 -0
  374. app/db/alembic/versions/20260806_000000_add_additional_usage_alias_probe_indexes.py +84 -0
  375. app/db/alembic/versions/20260806_000000_add_anonymous_telemetry.py +49 -0
  376. app/db/alembic/versions/20260806_010000_add_conversation_presence_rollup.py +72 -0
  377. app/db/alembic/versions/20260806_020000_add_usage_history_bulk_covering_indexes.py +99 -0
  378. app/db/alembic/versions/20260806_030000_add_api_key_allowed_reasoning_efforts.py +65 -0
  379. app/db/alembic/versions/20260806_120000_add_http_bridge_owner_process_epoch.py +57 -0
  380. app/db/alembic/versions/20260807_000000_merge_http_bridge_operation_spool_and_latest_main.py +24 -0
  381. app/db/alembic/versions/20260808_000000_tune_usage_history_autovacuum.py +54 -0
  382. app/db/alembic/versions/20260810_000000_add_http_bridge_recovery_dispatch_count.py +44 -0
  383. app/db/alembic/versions/20260811_000000_add_hourly_rollup_cancelled_count.py +92 -0
  384. app/db/alembic/versions/20260812_000000_merge_recovery_dispatch_and_hourly_cancelled_heads.py +25 -0
  385. app/db/alembic/versions/20260812_120000_add_sticky_abandonment_scope.py +43 -0
  386. app/db/alembic/versions/20260813_000000_add_file_account_pins.py +40 -0
  387. app/db/alembic/versions/20260816_000000_add_account_pending_deletion.py +81 -0
  388. app/db/alembic/versions/20260816_000000_add_model_source_embeddings.py +50 -0
  389. app/db/alembic/versions/20260821_000000_add_retry_circuit_admission_generation.py +49 -0
  390. app/db/alembic/versions/20260826_000000_add_http_bridge_event_chunks.py +107 -0
  391. app/db/alembic/versions/20260828_000000_add_accounts_chatgpt_identity_index.py +55 -0
  392. app/db/alembic/versions/20260830_000000_add_quota_warmup_claim_expiry.py +83 -0
  393. app/db/alembic/versions/20260908_000000_add_guest_session_generation.py +37 -0
  394. app/db/alembic/versions/20260908_000000_add_subscription_overflow.py +140 -0
  395. app/db/alembic/versions/20260908_000000_replace_upstream_stream_transport_default_sentinel.py +50 -0
  396. app/db/alembic/versions/20260908_020000_merge_overflow_transport_heads.py +27 -0
  397. app/db/alembic/versions/20260909_000000_add_dashboard_roles.py +72 -0
  398. app/db/alembic/versions/20260909_010000_add_dashboard_users.py +183 -0
  399. app/db/alembic/versions/20260909_020000_reproject_compat_admin_credentials.py +105 -0
  400. app/db/alembic/versions/20260909_030000_add_audit_actor_columns.py +103 -0
  401. app/db/alembic/versions/20260909_030000_dashboard_resilience_toggle_settings.py +63 -0
  402. app/db/alembic/versions/20260909_040000_add_dashboard_user_invites.py +54 -0
  403. app/db/alembic/versions/20260909_040000_dashboard_timeout_settings.py +68 -0
  404. app/db/alembic/versions/20260909_050000_dashboard_routing_overload_settings.py +65 -0
  405. app/db/alembic/versions/20260909_060000_add_report_rollup.py +74 -0
  406. app/db/alembic/versions/20260909_070000_automation_run_claim_budget.py +50 -0
  407. app/db/alembic/versions/20260909_080000_dashboard_stream_bridge_budgets.py +62 -0
  408. app/db/alembic/versions/20260909_090000_dashboard_background_job_toggles.py +63 -0
  409. app/db/alembic/versions/20260909_100000_dashboard_codex_prewarm.py +49 -0
  410. app/db/alembic/versions/20260909_110000_model_context_window_overrides.py +45 -0
  411. app/db/alembic/versions/20260909_120000_dashboard_conversation_archive.py +48 -0
  412. app/db/alembic/versions/20260909_130000_add_request_logs_live_facet_indexes.py +96 -0
  413. app/db/alembic/versions/20260910_000000_add_totp_required_for_admin_role.py +44 -0
  414. app/db/alembic/versions/20260910_000000_request_logs_missing_cost_index.py +52 -0
  415. app/db/alembic/versions/20260910_010000_add_dashboard_auth_providers.py +134 -0
  416. app/db/alembic/versions/20260910_010000_dashboard_spool_retention.py +52 -0
  417. app/db/alembic/versions/20260910_020000_add_dashboard_role_mappings.py +74 -0
  418. app/db/alembic/versions/20260910_143000_request_log_affinity.py +28 -0
  419. app/db/alembic/versions/20260910_180000_merge_affinity_guest_heads.py +17 -0
  420. app/db/alembic/versions/20260910_200000_merge_affinity_identity_heads.py +17 -0
  421. app/db/alembic/versions/20260910_220000_merge_affinity_invite_heads.py +17 -0
  422. app/db/alembic/versions/20260911_000000_model_source_pins_kind_expires_index.py +43 -0
  423. app/db/alembic/versions/20260911_010000_merge_pin_index_and_affinity_heads.py +17 -0
  424. app/db/alembic/versions/20260911_020000_add_http_bridge_terminal_append_phase.py +54 -0
  425. app/db/alembic/versions/20260911_030000_add_local_login_policy.py +48 -0
  426. app/db/alembic/versions/20260911_040000_add_thread_cache_identity_mode.py +48 -0
  427. app/db/alembic/versions/20260911_060000_add_bridge_session_continuity_abandonment.py +40 -0
  428. app/db/alembic/versions/20260912_000000_merge_thread_cache_and_bridge_retirement_heads.py +22 -0
  429. app/db/alembic/versions/20260912_010000_drop_legacy_dashboard_credentials.py +140 -0
  430. app/db/alembic/versions/20260913_000000_add_oidc_provider_flow.py +134 -0
  431. app/db/alembic/versions/20260914_000000_add_scim_tokens.py +78 -0
  432. app/db/alembic/versions/20260914_000000_drop_subscription_overflow_schema.py +132 -0
  433. app/db/alembic/versions/20260918_000000_merge_scim_and_overflow_heads.py +17 -0
  434. app/db/alembic/versions/__init__.py +1 -0
  435. app/db/backup.py +64 -0
  436. app/db/migrate.py +1249 -0
  437. app/db/migration_lock.py +160 -0
  438. app/db/migration_url.py +18 -0
  439. app/db/models.py +3100 -0
  440. app/db/recover.py +244 -0
  441. app/db/session.py +1295 -0
  442. app/db/snapshot.py +46 -0
  443. app/db/sqlite_lock_retry.py +190 -0
  444. app/db/sqlite_utils.py +549 -0
  445. app/dependencies.py +462 -0
  446. app/main.py +1240 -0
  447. app/modules/__init__.py +0 -0
  448. app/modules/accounts/__init__.py +0 -0
  449. app/modules/accounts/api.py +426 -0
  450. app/modules/accounts/auth_manager.py +1439 -0
  451. app/modules/accounts/background_repository.py +160 -0
  452. app/modules/accounts/deletion.py +540 -0
  453. app/modules/accounts/mappers.py +654 -0
  454. app/modules/accounts/refresh_claims.py +352 -0
  455. app/modules/accounts/repository.py +1679 -0
  456. app/modules/accounts/schemas.py +308 -0
  457. app/modules/accounts/service.py +809 -0
  458. app/modules/accounts/usage_rollup.py +486 -0
  459. app/modules/accounts/usage_rollup_scheduler.py +75 -0
  460. app/modules/accounts/usage_time_rollup.py +1037 -0
  461. app/modules/accounts/usage_time_rollup_read.py +522 -0
  462. app/modules/api_keys/__init__.py +1 -0
  463. app/modules/api_keys/api.py +344 -0
  464. app/modules/api_keys/last_used_coalescer.py +216 -0
  465. app/modules/api_keys/limit_windows.py +41 -0
  466. app/modules/api_keys/repository.py +1200 -0
  467. app/modules/api_keys/reset_scheduler.py +85 -0
  468. app/modules/api_keys/schemas.py +135 -0
  469. app/modules/api_keys/service.py +2123 -0
  470. app/modules/audit/__init__.py +0 -0
  471. app/modules/audit/api.py +88 -0
  472. app/modules/audit/repository.py +74 -0
  473. app/modules/audit/schemas.py +33 -0
  474. app/modules/audit/service.py +71 -0
  475. app/modules/auth_providers/__init__.py +0 -0
  476. app/modules/auth_providers/api.py +127 -0
  477. app/modules/auth_providers/config.py +343 -0
  478. app/modules/auth_providers/repository.py +101 -0
  479. app/modules/auth_providers/schemas.py +88 -0
  480. app/modules/auth_providers/seed.py +105 -0
  481. app/modules/auth_providers/service.py +289 -0
  482. app/modules/automations/__init__.py +1 -0
  483. app/modules/automations/api.py +369 -0
  484. app/modules/automations/repository.py +1896 -0
  485. app/modules/automations/scheduler.py +107 -0
  486. app/modules/automations/schemas.py +158 -0
  487. app/modules/automations/service.py +2548 -0
  488. app/modules/cache_isolation_probe/__init__.py +0 -0
  489. app/modules/cache_isolation_probe/api.py +188 -0
  490. app/modules/cache_isolation_probe/prefix.py +144 -0
  491. app/modules/cache_isolation_probe/schemas.py +94 -0
  492. app/modules/cache_isolation_probe/sender.py +229 -0
  493. app/modules/cache_isolation_probe/service.py +525 -0
  494. app/modules/conversation_archive/__init__.py +1 -0
  495. app/modules/conversation_archive/api.py +101 -0
  496. app/modules/conversation_archive/schemas.py +28 -0
  497. app/modules/conversation_archive/service.py +165 -0
  498. app/modules/dashboard/__init__.py +1 -0
  499. app/modules/dashboard/api.py +106 -0
  500. app/modules/dashboard/builders.py +107 -0
  501. app/modules/dashboard/repository.py +304 -0
  502. app/modules/dashboard/schemas.py +148 -0
  503. app/modules/dashboard/service.py +534 -0
  504. app/modules/dashboard/timeframes.py +19 -0
  505. app/modules/dashboard/weekly_pace.py +617 -0
  506. app/modules/dashboard_auth/__init__.py +1 -0
  507. app/modules/dashboard_auth/api.py +1344 -0
  508. app/modules/dashboard_auth/oidc_api.py +753 -0
  509. app/modules/dashboard_auth/oidc_flows.py +330 -0
  510. app/modules/dashboard_auth/repository.py +452 -0
  511. app/modules/dashboard_auth/schemas.py +232 -0
  512. app/modules/dashboard_auth/service.py +1384 -0
  513. app/modules/dashboard_roles/__init__.py +0 -0
  514. app/modules/dashboard_roles/api.py +64 -0
  515. app/modules/dashboard_roles/repository.py +54 -0
  516. app/modules/dashboard_roles/schemas.py +29 -0
  517. app/modules/dashboard_roles/seed.py +80 -0
  518. app/modules/dashboard_roles/service.py +142 -0
  519. app/modules/dashboard_users/__init__.py +0 -0
  520. app/modules/dashboard_users/api.py +278 -0
  521. app/modules/dashboard_users/break_glass.py +169 -0
  522. app/modules/dashboard_users/credentials.py +18 -0
  523. app/modules/dashboard_users/identity_resolver.py +665 -0
  524. app/modules/dashboard_users/repository.py +725 -0
  525. app/modules/dashboard_users/schemas.py +117 -0
  526. app/modules/dashboard_users/service.py +1054 -0
  527. app/modules/firewall/__init__.py +1 -0
  528. app/modules/firewall/api.py +78 -0
  529. app/modules/firewall/repository.py +51 -0
  530. app/modules/firewall/schemas.py +28 -0
  531. app/modules/firewall/service.py +93 -0
  532. app/modules/fleet/__init__.py +1 -0
  533. app/modules/fleet/api.py +177 -0
  534. app/modules/fleet/mappers.py +56 -0
  535. app/modules/fleet/observability.py +372 -0
  536. app/modules/fleet/schemas.py +119 -0
  537. app/modules/health/__init__.py +0 -0
  538. app/modules/health/api.py +246 -0
  539. app/modules/health/schemas.py +27 -0
  540. app/modules/limit_warmup/__init__.py +1 -0
  541. app/modules/limit_warmup/repository.py +159 -0
  542. app/modules/limit_warmup/service.py +1001 -0
  543. app/modules/model_sources/__init__.py +1 -0
  544. app/modules/model_sources/api.py +99 -0
  545. app/modules/model_sources/catalog.py +307 -0
  546. app/modules/model_sources/forwarding.py +1435 -0
  547. app/modules/model_sources/projection.py +55 -0
  548. app/modules/model_sources/repository.py +201 -0
  549. app/modules/model_sources/schemas.py +79 -0
  550. app/modules/model_sources/selection.py +166 -0
  551. app/modules/model_sources/service.py +236 -0
  552. app/modules/oauth/__init__.py +0 -0
  553. app/modules/oauth/api.py +101 -0
  554. app/modules/oauth/repository.py +238 -0
  555. app/modules/oauth/schemas.py +49 -0
  556. app/modules/oauth/service.py +1196 -0
  557. app/modules/oauth/templates/oauth_success.html +122 -0
  558. app/modules/proxy/__init__.py +0 -0
  559. app/modules/proxy/_load_balancer/__init__.py +1 -0
  560. app/modules/proxy/_load_balancer/error_rate.py +116 -0
  561. app/modules/proxy/_load_balancer/latency_cohort.py +154 -0
  562. app/modules/proxy/_load_balancer/model_eligibility.py +250 -0
  563. app/modules/proxy/_load_balancer/opportunistic_admission.py +350 -0
  564. app/modules/proxy/_load_balancer/overload_backoff.py +394 -0
  565. app/modules/proxy/_load_balancer/sticky_selection.py +2417 -0
  566. app/modules/proxy/_load_balancer/throughput_cohort.py +233 -0
  567. app/modules/proxy/_load_balancer/ttft_cohort.py +214 -0
  568. app/modules/proxy/_load_balancer/tunables.py +97 -0
  569. app/modules/proxy/_load_balancer/types.py +96 -0
  570. app/modules/proxy/_load_balancer/unbound_selection.py +534 -0
  571. app/modules/proxy/_service/__init__.py +6 -0
  572. app/modules/proxy/_service/api_key_usage.py +745 -0
  573. app/modules/proxy/_service/clock_budget.py +38 -0
  574. app/modules/proxy/_service/codex_control.py +611 -0
  575. app/modules/proxy/_service/compact.py +2211 -0
  576. app/modules/proxy/_service/file_ops.py +749 -0
  577. app/modules/proxy/_service/http_bridge/__init__.py +3 -0
  578. app/modules/proxy/_service/http_bridge/accepted_replay.py +331 -0
  579. app/modules/proxy/_service/http_bridge/account_sessions.py +57 -0
  580. app/modules/proxy/_service/http_bridge/activity.py +178 -0
  581. app/modules/proxy/_service/http_bridge/helpers.py +3959 -0
  582. app/modules/proxy/_service/http_bridge/mixin.py +2435 -0
  583. app/modules/proxy/_service/http_bridge/owner_forwarding.py +609 -0
  584. app/modules/proxy/_service/http_bridge/protocol.py +167 -0
  585. app/modules/proxy/_service/http_bridge/proxy_failover.py +58 -0
  586. app/modules/proxy/_service/http_bridge/quarantine.py +443 -0
  587. app/modules/proxy/_service/http_bridge/request_submit.py +4582 -0
  588. app/modules/proxy/_service/http_bridge/retry_circuit.py +2442 -0
  589. app/modules/proxy/_service/http_bridge/service_stubs.py +492 -0
  590. app/modules/proxy/_service/http_bridge/session_registry.py +778 -0
  591. app/modules/proxy/_service/http_bridge/streaming.py +5328 -0
  592. app/modules/proxy/_service/http_bridge/upstream_events.py +4538 -0
  593. app/modules/proxy/_service/observability.py +410 -0
  594. app/modules/proxy/_service/rate_limit.py +361 -0
  595. app/modules/proxy/_service/realtime_live.py +655 -0
  596. app/modules/proxy/_service/refresh.py +154 -0
  597. app/modules/proxy/_service/request_log.py +585 -0
  598. app/modules/proxy/_service/response_create.py +930 -0
  599. app/modules/proxy/_service/streaming/__init__.py +3 -0
  600. app/modules/proxy/_service/streaming/helpers.py +1317 -0
  601. app/modules/proxy/_service/streaming/mixin.py +1100 -0
  602. app/modules/proxy/_service/streaming/protocol.py +32 -0
  603. app/modules/proxy/_service/streaming/retry.py +3762 -0
  604. app/modules/proxy/_service/support.py +2122 -0
  605. app/modules/proxy/_service/transcribe.py +434 -0
  606. app/modules/proxy/_service/warmup.py +483 -0
  607. app/modules/proxy/_service/websocket/__init__.py +3 -0
  608. app/modules/proxy/_service/websocket/helpers.py +2360 -0
  609. app/modules/proxy/_service/websocket/mixin.py +7285 -0
  610. app/modules/proxy/_service/websocket/protocol.py +74 -0
  611. app/modules/proxy/account_cache.py +229 -0
  612. app/modules/proxy/account_eligibility.py +87 -0
  613. app/modules/proxy/additional_model_limits.py +40 -0
  614. app/modules/proxy/affinity.py +830 -0
  615. app/modules/proxy/affinity_observation.py +101 -0
  616. app/modules/proxy/api.py +10254 -0
  617. app/modules/proxy/api_key_usage.py +113 -0
  618. app/modules/proxy/cap_partitioning.py +280 -0
  619. app/modules/proxy/capability_lineage.py +65 -0
  620. app/modules/proxy/capability_lineage_repository.py +82 -0
  621. app/modules/proxy/capability_routing.py +253 -0
  622. app/modules/proxy/continuity.py +112 -0
  623. app/modules/proxy/downstream_delivery.py +277 -0
  624. app/modules/proxy/durable_bridge_coordinator.py +1217 -0
  625. app/modules/proxy/durable_bridge_repository.py +4516 -0
  626. app/modules/proxy/durable_bridge_runtime.py +9 -0
  627. app/modules/proxy/durable_bridge_transcript_codec.py +122 -0
  628. app/modules/proxy/fair_share.py +131 -0
  629. app/modules/proxy/file_pin_repository.py +210 -0
  630. app/modules/proxy/helpers.py +439 -0
  631. app/modules/proxy/http_bridge_event_batcher.py +613 -0
  632. app/modules/proxy/http_bridge_forwarding.py +797 -0
  633. app/modules/proxy/http_continuation.py +76 -0
  634. app/modules/proxy/images_observability.py +62 -0
  635. app/modules/proxy/images_service.py +1054 -0
  636. app/modules/proxy/load_balancer.py +3016 -0
  637. app/modules/proxy/rate_limit_cache.py +46 -0
  638. app/modules/proxy/replay_safety.py +1063 -0
  639. app/modules/proxy/repo_bundle.py +31 -0
  640. app/modules/proxy/request_policy.py +1054 -0
  641. app/modules/proxy/ring_membership.py +251 -0
  642. app/modules/proxy/schemas.py +345 -0
  643. app/modules/proxy/selection_errors.py +41 -0
  644. app/modules/proxy/service.py +2600 -0
  645. app/modules/proxy/source_admission.py +125 -0
  646. app/modules/proxy/source_dispatch.py +999 -0
  647. app/modules/proxy/sticky_repository.py +831 -0
  648. app/modules/proxy/thread_anchors.py +491 -0
  649. app/modules/proxy/tool_call_dedupe.py +802 -0
  650. app/modules/proxy/types.py +54 -0
  651. app/modules/proxy/work_admission.py +121 -0
  652. app/modules/quota_planner/__init__.py +1 -0
  653. app/modules/quota_planner/api.py +276 -0
  654. app/modules/quota_planner/logic.py +743 -0
  655. app/modules/quota_planner/repository.py +722 -0
  656. app/modules/quota_planner/scheduler.py +212 -0
  657. app/modules/quota_planner/schemas.py +102 -0
  658. app/modules/quota_planner/warmup.py +693 -0
  659. app/modules/rate_limit_reset_credits/__init__.py +0 -0
  660. app/modules/rate_limit_reset_credits/api.py +601 -0
  661. app/modules/rate_limit_reset_credits/redeem_coordination.py +274 -0
  662. app/modules/rate_limit_reset_credits/store.py +114 -0
  663. app/modules/reports/__init__.py +1 -0
  664. app/modules/reports/api.py +134 -0
  665. app/modules/reports/cache.py +74 -0
  666. app/modules/reports/filters.py +27 -0
  667. app/modules/reports/repository.py +525 -0
  668. app/modules/reports/rollup.py +150 -0
  669. app/modules/reports/rollup_read.py +148 -0
  670. app/modules/reports/schemas.py +123 -0
  671. app/modules/reports/service.py +287 -0
  672. app/modules/reports/thread_identity.py +256 -0
  673. app/modules/request_logs/__init__.py +0 -0
  674. app/modules/request_logs/api.py +224 -0
  675. app/modules/request_logs/cost_backfill.py +154 -0
  676. app/modules/request_logs/mappers.py +97 -0
  677. app/modules/request_logs/repository.py +1765 -0
  678. app/modules/request_logs/schemas.py +145 -0
  679. app/modules/request_logs/service.py +348 -0
  680. app/modules/role_mappings/__init__.py +0 -0
  681. app/modules/role_mappings/api.py +206 -0
  682. app/modules/role_mappings/matching.py +55 -0
  683. app/modules/role_mappings/repository.py +106 -0
  684. app/modules/role_mappings/schemas.py +67 -0
  685. app/modules/role_mappings/service.py +334 -0
  686. app/modules/runtime/__init__.py +1 -0
  687. app/modules/runtime/api.py +18 -0
  688. app/modules/runtime/schemas.py +14 -0
  689. app/modules/runtime/service.py +218 -0
  690. app/modules/scim/__init__.py +1 -0
  691. app/modules/scim/api.py +194 -0
  692. app/modules/scim/dependencies.py +112 -0
  693. app/modules/scim/errors.py +77 -0
  694. app/modules/scim/management_api.py +156 -0
  695. app/modules/scim/repository.py +212 -0
  696. app/modules/scim/schemas.py +145 -0
  697. app/modules/scim/service.py +786 -0
  698. app/modules/scim/tokens.py +43 -0
  699. app/modules/settings/__init__.py +1 -0
  700. app/modules/settings/api.py +1677 -0
  701. app/modules/settings/repository.py +606 -0
  702. app/modules/settings/schemas.py +496 -0
  703. app/modules/settings/service.py +769 -0
  704. app/modules/shared/__init__.py +0 -0
  705. app/modules/shared/schemas.py +22 -0
  706. app/modules/sticky_sessions/__init__.py +1 -0
  707. app/modules/sticky_sessions/api.py +118 -0
  708. app/modules/sticky_sessions/cleanup_scheduler.py +487 -0
  709. app/modules/sticky_sessions/schemas.py +78 -0
  710. app/modules/sticky_sessions/service.py +196 -0
  711. app/modules/telemetry/__init__.py +5 -0
  712. app/modules/telemetry/api.py +115 -0
  713. app/modules/telemetry/clients.py +62 -0
  714. app/modules/telemetry/consent.py +115 -0
  715. app/modules/telemetry/scheduler.py +85 -0
  716. app/modules/telemetry/schemas.py +163 -0
  717. app/modules/telemetry/sender.py +221 -0
  718. app/modules/telemetry/snapshot.py +490 -0
  719. app/modules/usage/__init__.py +0 -0
  720. app/modules/usage/additional_quota_keys.py +379 -0
  721. app/modules/usage/api.py +45 -0
  722. app/modules/usage/background_repository.py +147 -0
  723. app/modules/usage/builders.py +559 -0
  724. app/modules/usage/depletion_service.py +277 -0
  725. app/modules/usage/live_ingest.py +356 -0
  726. app/modules/usage/mappers.py +29 -0
  727. app/modules/usage/plan_downgrade_observations.py +488 -0
  728. app/modules/usage/repository.py +1890 -0
  729. app/modules/usage/schemas.py +68 -0
  730. app/modules/usage/service.py +102 -0
  731. app/modules/usage/updater.py +1488 -0
  732. app/static/assets/Trans-Cj-J9mVA.js +1 -0
  733. app/static/assets/access-page-DHzZiXZi.js +1 -0
  734. app/static/assets/account-cost-donut-DGhfA4U8.js +1 -0
  735. app/static/assets/account-identifiers-CxbdLshZ.js +1 -0
  736. app/static/assets/account-status-DIsU5sUS.js +1 -0
  737. app/static/assets/account-trend-chart-BptlfwwQ.js +1 -0
  738. app/static/assets/accounts-page-Bbai8ift.js +3 -0
  739. app/static/assets/api-Dwa7XxPe.js +1 -0
  740. app/static/assets/api-client-dugX1Pti.js +55 -0
  741. app/static/assets/api-key-create-dialog-DnRm0oYN.js +1 -0
  742. app/static/assets/api-key-created-dialog-B6jkDG9f.js +1 -0
  743. app/static/assets/api-key-edit-dialog-Ck3Bnc3V.js +1 -0
  744. app/static/assets/api-trend-chart-BchCw7Hj.js +1 -0
  745. app/static/assets/apis-page-BvXz9c6i.js +2 -0
  746. app/static/assets/arrow-up-Dqh_KzJo.js +1 -0
  747. app/static/assets/arrow-up-down-LBRgVDGY.js +1 -0
  748. app/static/assets/automations-page-DQgNAJOR.js +3 -0
  749. app/static/assets/badge-DqMb9es0.js +1 -0
  750. app/static/assets/blur-email-HEvKagG2.js +1 -0
  751. app/static/assets/button-BjVmmiGW.js +1 -0
  752. app/static/assets/checkbox-apsl9vzk.js +1 -0
  753. app/static/assets/chevron-left-1d_pHVa5.js +1 -0
  754. app/static/assets/circle-alert-BS0xd_o1.js +1 -0
  755. app/static/assets/clsx-DB0hHKMi.js +1 -0
  756. app/static/assets/colors-BBy8GEtr.js +1 -0
  757. app/static/assets/copy-button-D7nAvOKR.js +1 -0
  758. app/static/assets/cost-per-day-chart-YdqlkZVM.js +1 -0
  759. app/static/assets/createLucideIcon-DBLzLNh2.js +1 -0
  760. app/static/assets/dashboard-page-CSdhBy0f.js +2 -0
  761. app/static/assets/distribution-metric-format-yZTy1ZEu.js +1 -0
  762. app/static/assets/donut-chart-Bxc69cra.js +1 -0
  763. app/static/assets/download-kQteIcsB.js +1 -0
  764. app/static/assets/ellipsis-B_5OgIre.js +1 -0
  765. app/static/assets/empty-state-JcA7tc8m.js +1 -0
  766. app/static/assets/es2015-bMH8YKZT.js +41 -0
  767. app/static/assets/es6-B264ni6W.js +69 -0
  768. app/static/assets/external-link-UVAOgMQ0.js +1 -0
  769. app/static/assets/i18n-BWqZ9Ua2.js +1 -0
  770. app/static/assets/inbox-B_RM2s1J.js +1 -0
  771. app/static/assets/index-DIpd264Z.js +23 -0
  772. app/static/assets/index-DMeLKRJx.css +2 -0
  773. app/static/assets/inherit-badge-QhTYVaU6.js +1 -0
  774. app/static/assets/invite-dialog-0CYh0GiQ.js +1 -0
  775. app/static/assets/iso-Bpe5Ihbb.js +1 -0
  776. app/static/assets/label-CGaLmLNf.js +1 -0
  777. app/static/assets/lazy-recharts-Pv8p9ArM.js +2 -0
  778. app/static/assets/loading-overlay-CzKQfHLI.js +1 -0
  779. app/static/assets/mini-quota-bar-cGT_t8Ae.js +1 -0
  780. app/static/assets/model-distribution-donut-BKjAfT48.js +1 -0
  781. app/static/assets/model-source-multi-select-D2d06Vc7.js +1 -0
  782. app/static/assets/multi-select-filter-KIA5qU6Z.js +1 -0
  783. app/static/assets/oauth-dialog-XeaGGDjS.js +1 -0
  784. app/static/assets/pagination-controls-D1Ypd0RU.js +1 -0
  785. app/static/assets/payload-CIYnGZz9.js +1 -0
  786. app/static/assets/queue-wait-chart-VoV7Yg9L.js +1 -0
  787. app/static/assets/report-chart-card-BXDliwcc.js +1 -0
  788. app/static/assets/reports-page-CGCJeMHG.js +3 -0
  789. app/static/assets/reset-credit-confirm-dialog-QOifTnVB.js +1 -0
  790. app/static/assets/rolldown-runtime-hePW80VL.js +1 -0
  791. app/static/assets/search-HfhqwERt.js +1 -0
  792. app/static/assets/select-Cl33bz-g.js +1 -0
  793. app/static/assets/settings-page-mZj01HYl.js +2 -0
  794. app/static/assets/shim-CZwyM4Q7.js +1 -0
  795. app/static/assets/skeleton-C_EUc4Rp.js +1 -0
  796. app/static/assets/sparkline-chart-DdNNsddM.js +1 -0
  797. app/static/assets/status-badge-CHHA5X6c.js +1 -0
  798. app/static/assets/switch-CzxgokcM.js +1 -0
  799. app/static/assets/table-DXaNSQ0G.js +1 -0
  800. app/static/assets/time-to-first-token-chart--68L0u_v.js +1 -0
  801. app/static/assets/tokens-per-day-chart-Cs82BFuE.js +1 -0
  802. app/static/assets/tokens-per-second-chart-AD2dfOn6.js +1 -0
  803. app/static/assets/totp-settings-CRHtxQ7j.js +1 -0
  804. app/static/assets/trash-D3KyWr1h.js +1 -0
  805. app/static/assets/use-account-quota-display-BPmrTBTC.js +1 -0
  806. app/static/assets/use-accounts-PWlVGw3z.js +1 -0
  807. app/static/assets/use-auth-CuDlW3iC.js +1 -0
  808. app/static/assets/use-chart-colors-DIu6tClI.js +1 -0
  809. app/static/assets/use-model-sources-Bul24JBu.js +1 -0
  810. app/static/assets/use-models-DUCPUuZg.js +1 -0
  811. app/static/assets/use-reduced-motion-hAm5r6_d.js +1 -0
  812. app/static/assets/useTranslation--56nHhlH.js +1 -0
  813. app/static/assets/useragent-distribution-donut-Dup6sf_n.js +1 -0
  814. app/static/assets/users-uIC3TIlE.js +1 -0
  815. app/static/assets/vendor-query-_bHZuxEx.js +1 -0
  816. app/static/assets/vendor-react-BVL0YZPW.js +9 -0
  817. app/static/favicon.svg +10 -0
  818. app/static/fonts/GeistSans-Variable.woff2 +0 -0
  819. app/static/fonts/jetbrains-mono-latin-ext.woff2 +0 -0
  820. app/static/fonts/jetbrains-mono-latin.woff2 +0 -0
  821. app/static/index.html +29 -0
  822. codexlb_up-1.25.0b10.dist-info/METADATA +124 -0
  823. codexlb_up-1.25.0b10.dist-info/RECORD +827 -0
  824. codexlb_up-1.25.0b10.dist-info/WHEEL +4 -0
  825. codexlb_up-1.25.0b10.dist-info/entry_points.txt +3 -0
  826. codexlb_up-1.25.0b10.dist-info/licenses/LICENSE +21 -0
  827. config/additional_quota_registry.json +11 -0
@@ -0,0 +1,3016 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ import inspect
5
+ import json
6
+ import logging
7
+ from collections.abc import Collection, Mapping
8
+ from dataclasses import dataclass, replace
9
+ from datetime import datetime, timedelta, timezone
10
+ from typing import TYPE_CHECKING, Iterable
11
+ from uuid import uuid4
12
+
13
+ from app.core import usage as usage_core
14
+ from app.core.balancer import (
15
+ ERROR_BACKOFF_THRESHOLD,
16
+ HEALTH_TIER_DRAINING,
17
+ HEALTH_TIER_HEALTHY,
18
+ HEALTH_TIER_PROBING,
19
+ QUOTA_EXCEEDED_COOLDOWN_SECONDS,
20
+ RATE_LIMITED_MIN_COOLDOWN_SECONDS,
21
+ ROUTING_POLICY_BURN_FIRST,
22
+ ROUTING_POLICY_PRESERVE,
23
+ TRAFFIC_CLASS_FOREGROUND,
24
+ TRAFFIC_CLASS_OPPORTUNISTIC,
25
+ AccountState,
26
+ ResetPreferenceWindow,
27
+ RoutingCostsByAccount,
28
+ RoutingStrategy,
29
+ TrafficClass,
30
+ evaluate_health_tier,
31
+ handle_permanent_failure,
32
+ handle_quota_exceeded,
33
+ handle_rate_limit,
34
+ plausible_rate_limit_reset_at,
35
+ )
36
+ from app.core.balancer import (
37
+ select_account as select_account,
38
+ )
39
+ from app.core.balancer.types import UpstreamError
40
+ from app.core.clock import REAL_CLOCK, Clock
41
+ from app.core.config.dashboard_overrides import with_dashboard_overrides
42
+ from app.core.config.settings import get_settings
43
+ from app.core.config.settings_cache import get_settings_cache
44
+ from app.core.crypto import TokenEncryptor
45
+ from app.core.metrics.prometheus import (
46
+ PROMETHEUS_AVAILABLE,
47
+ account_cap_rejections_total,
48
+ account_inflight_leases,
49
+ account_lease_acquired_total,
50
+ account_lease_released_total,
51
+ account_lease_stale_reclaimed_total,
52
+ api_key_fair_share_rejections_total,
53
+ stream_pool_capacity,
54
+ stream_pool_inflight,
55
+ )
56
+ from app.core.openai.model_registry import get_model_registry
57
+ from app.core.plan_types import account_plan_matches_allowed, normalize_account_plan_type
58
+ from app.core.resilience.circuit_breaker import are_all_account_circuit_breakers_open
59
+ from app.core.resilience.degradation import get_status as get_degradation_status
60
+ from app.core.resilience.degradation import set_degraded, set_normal
61
+ from app.core.resilience.toggles import resolve_resilience_toggles
62
+ from app.core.usage.quota import apply_usage_quota
63
+ from app.core.usage.refresh_policy import usage_freshness_horizon_seconds
64
+ from app.core.utils.time import to_utc_naive, utcnow
65
+ from app.db.models import Account, AccountStatus, AdditionalUsageHistory, StickySessionKind, UsageHistory
66
+ from app.db.snapshot import clone_row
67
+ from app.modules.proxy._load_balancer.error_rate import (
68
+ ErrorRateWeightingPolicy,
69
+ error_rate_weight_multiplier,
70
+ record_outcome_locked,
71
+ )
72
+ from app.modules.proxy._load_balancer.latency_cohort import apply_latency_cohort_weights
73
+ from app.modules.proxy._load_balancer.model_eligibility import (
74
+ _ADDITIONAL_QUOTA_EXEMPT_PLAN_TYPES,
75
+ CatalogOmissionQuotaAdmission,
76
+ _additional_quota_applies_to_plan, # noqa: F401
77
+ _additional_quota_eligibility,
78
+ _additional_usage_is_exhausted, # noqa: F401
79
+ _catalog_omission_quota_admission,
80
+ _effective_model_service_tier, # noqa: F401
81
+ _gated_limit_name_for_model,
82
+ _latest_additional_by_key,
83
+ _ModelAccountFilterResult,
84
+ _normalize_model_id, # noqa: F401
85
+ )
86
+ from app.modules.proxy._load_balancer.model_eligibility import (
87
+ _filter_accounts_for_model as _filter_accounts_for_model_impl,
88
+ )
89
+ from app.modules.proxy._load_balancer.model_eligibility import (
90
+ _filter_accounts_for_model_with_catalog_evidence as _filter_accounts_for_model_with_catalog_evidence_impl,
91
+ )
92
+ from app.modules.proxy._load_balancer.model_eligibility import (
93
+ _mapped_model_has_registry_entry as _mapped_model_has_registry_entry_impl,
94
+ )
95
+ from app.modules.proxy._load_balancer.opportunistic_admission import (
96
+ OPPORTUNISTIC_BURN_WINDOW_CLOSED,
97
+ OpportunisticAdmissionRequest,
98
+ apply_lease_release,
99
+ detached_runtime_snapshot,
100
+ run_opportunistic_admission,
101
+ )
102
+ from app.modules.proxy._load_balancer.sticky_selection import (
103
+ _STICKY_EXISTING_UNSET,
104
+ SelectionInputsProtocol,
105
+ StickySelectionRequest,
106
+ _clone_account,
107
+ _StickySelectionOutcome,
108
+ prepare_selection_states,
109
+ run_sticky_selection_path,
110
+ )
111
+ from app.modules.proxy._load_balancer.sticky_selection import (
112
+ _account_cap_error_message as _account_cap_error_message,
113
+ )
114
+ from app.modules.proxy._load_balancer.sticky_selection import (
115
+ _best_health_tier_states as _best_health_tier_states,
116
+ )
117
+ from app.modules.proxy._load_balancer.sticky_selection import (
118
+ _filter_recovery_probe_candidates as _filter_recovery_probe_candidates,
119
+ )
120
+ from app.modules.proxy._load_balancer.sticky_selection import _persist_sticky_mutation as _persist_sticky_mutation
121
+ from app.modules.proxy._load_balancer.sticky_selection import (
122
+ _probing_result_requires_recovery_reservation as _probing_result_requires_recovery_reservation,
123
+ )
124
+ from app.modules.proxy._load_balancer.sticky_selection import (
125
+ _restore_sticky_mutation as _restore_sticky_mutation,
126
+ )
127
+ from app.modules.proxy._load_balancer.sticky_selection import (
128
+ _select_account_preferring_budget_safe as _select_account_preferring_budget_safe,
129
+ )
130
+ from app.modules.proxy._load_balancer.sticky_selection import _select_with_stickiness as _run_select_with_stickiness
131
+ from app.modules.proxy._load_balancer.sticky_selection import (
132
+ _state_above_budget_threshold as _state_above_budget_threshold,
133
+ )
134
+ from app.modules.proxy._load_balancer.sticky_selection import (
135
+ _state_above_sticky_budget_threshold as _state_above_sticky_budget_threshold,
136
+ )
137
+ from app.modules.proxy._load_balancer.tunables import (
138
+ RoutingTunables,
139
+ account_lease_stale_ttl_seconds,
140
+ resolve_routing_tunables,
141
+ )
142
+ from app.modules.proxy._load_balancer.types import (
143
+ AccountConcurrencyCaps,
144
+ AccountLease,
145
+ AccountLeaseKind,
146
+ ProbeReservation,
147
+ RuntimeState,
148
+ )
149
+ from app.modules.proxy._load_balancer.unbound_selection import (
150
+ UnboundSelectionRequest,
151
+ run_unbound_selection_path,
152
+ )
153
+ from app.modules.proxy.account_cache import get_account_selection_cache, mark_account_routing_unavailable
154
+ from app.modules.proxy.account_eligibility import (
155
+ account_access_token_expires_at,
156
+ all_accounts_require_reauthentication,
157
+ )
158
+ from app.modules.proxy.affinity import _CodexSessionSource
159
+ from app.modules.proxy.cap_partitioning import (
160
+ configured_account_concurrency_caps,
161
+ get_cap_partition,
162
+ partition_cap,
163
+ )
164
+ from app.modules.proxy.fair_share import (
165
+ ApiKeyFairShareDenialError,
166
+ FairShareDecision,
167
+ effective_stream_pool_capacity,
168
+ evaluate_stream_fair_share,
169
+ )
170
+ from app.modules.proxy.repo_bundle import ProxyRepoFactory, ProxyRepositories
171
+ from app.modules.quota_planner.logic import PlannerSettings
172
+ from app.modules.usage.additional_quota_keys import (
173
+ canonicalize_additional_quota_key,
174
+ get_additional_quota_routing_policy,
175
+ normalize_additional_quota_key,
176
+ )
177
+ from app.modules.usage.mappers import usage_history_to_window_row
178
+
179
+ if TYPE_CHECKING:
180
+ from app.modules.accounts.repository import AccountsRepository
181
+ from app.modules.proxy.sticky_repository import StickyOwnerLookup, StickySessionsRepository
182
+
183
+ logger = logging.getLogger(__name__)
184
+
185
+ # Rows written by the same upstream fetch land within milliseconds of each
186
+ # other; a sibling row only proves a *later* fetch (one that no longer
187
+ # reported the stale window) when it is newer by more than this margin.
188
+ _SIBLING_FETCH_MARGIN_SECONDS = 5.0
189
+
190
+ _UsageWindowEntry = UsageHistory | AdditionalUsageHistory
191
+
192
+ NO_PLAN_SUPPORT_FOR_MODEL = "no_plan_support_for_model"
193
+ ADDITIONAL_QUOTA_DATA_UNAVAILABLE = "additional_quota_data_unavailable"
194
+ ADDITIONAL_QUOTA_EXHAUSTED = "quota_exhausted"
195
+ NO_ADDITIONAL_QUOTA_ELIGIBLE_ACCOUNTS = "no_additional_quota_eligible_accounts"
196
+ _ROUTING_POLICY_NORMAL = "normal"
197
+ _ACCOUNT_ROUTING_POLICIES = frozenset({_ROUTING_POLICY_NORMAL, ROUTING_POLICY_BURN_FIRST, ROUTING_POLICY_PRESERVE})
198
+ _ADDITIONAL_QUOTA_ROUTING_POLICIES = _ACCOUNT_ROUTING_POLICIES | frozenset({"inherit"})
199
+ CONTINUITY_OWNER_UNAVAILABLE = "continuity_owner_unavailable"
200
+ CONTINUITY_OWNER_POLICY_CONFLICT = "continuity_owner_policy_conflict"
201
+ _AMBIGUOUS_CONVERSATION_OWNER_CODE = "conversation_owner_unavailable"
202
+ _AMBIGUOUS_CONVERSATION_OWNER_MESSAGE = "Conversation owner cannot be determined from the eligible account pool"
203
+
204
+
205
+ @dataclass(frozen=True, slots=True)
206
+ class _NormalizedUsageInputs:
207
+ primary_used: float | None
208
+ primary_reset: int | None
209
+ primary_window_minutes: int | None
210
+ effective_secondary_entry: _UsageWindowEntry | None
211
+ secondary_used: float | None
212
+ secondary_reset: int | None
213
+
214
+
215
+ @dataclass
216
+ class AccountSelection:
217
+ account: Account | None
218
+ error_message: str | None
219
+ error_code: str | None = None
220
+ resets_at: int | None = None
221
+ lease: AccountLease | None = None
222
+ catalog_omission_quota_admission: CatalogOmissionQuotaAdmission | None = None
223
+ continuity_owner_no_longer_exists: bool = False
224
+ # ``hard_affinity_saturated`` whose resolved owner is one of the caller's
225
+ # own ``exclude_account_ids``: the wait a transient owner outage earns
226
+ # cannot clear this one (``_hard_affinity_owner_excluded_by_caller``).
227
+ hard_affinity_owner_excluded: bool = False
228
+
229
+
230
+ @dataclass(frozen=True, slots=True)
231
+ class _AdditionalLimitFilterResult:
232
+ accounts: list[Account]
233
+ latest_primary: dict[str, AdditionalUsageHistory]
234
+ latest_secondary: dict[str, AdditionalUsageHistory]
235
+ error_code: str | None = None
236
+ error_message: str | None = None
237
+
238
+
239
+ @dataclass(frozen=True, slots=True)
240
+ class _SelectionInputs(SelectionInputsProtocol):
241
+ accounts: list[Account]
242
+ latest_primary: dict[str, UsageHistory | AdditionalUsageHistory]
243
+ latest_secondary: dict[str, UsageHistory | AdditionalUsageHistory]
244
+ latest_monthly: dict[str, UsageHistory]
245
+ # Resolve ownership before transient routing filters; keep that stronger
246
+ # candidate pool alongside the effective routing pool.
247
+ continuity_owner_candidates: list[Account] | None = None
248
+ # Sticky mutation authority precedes model/service-tier eligibility; keep
249
+ # it separate because a model-ineligible account can still own the raw row
250
+ # this authenticated request may retire.
251
+ sticky_mutation_authority_account_ids: frozenset[str] | None = None
252
+ quota_planner_settings: PlannerSettings = PlannerSettings()
253
+ runtime_accounts: list[Account] | None = None
254
+ error_message: str | None = None
255
+ error_code: str | None = None
256
+ ignore_standard_quota_account_ids: frozenset[str] = frozenset()
257
+ ignore_standard_quota_status: bool = False
258
+ persist_standard_quota_status: bool = True
259
+ routing_policy_override: str | None = None
260
+ quota_admitted_catalog_omission_account_ids: frozenset[str] = frozenset()
261
+ # C2-3 resilience toggles: resolved once per selection from the dashboard
262
+ # snapshot passed to ``select_account``; None = inherit the env alias.
263
+ soft_drain_enabled: bool | None = None
264
+
265
+ @property
266
+ def effective_continuity_owner_candidates(self) -> list[Account]:
267
+ if self.continuity_owner_candidates is None:
268
+ return self.accounts
269
+ return self.continuity_owner_candidates
270
+
271
+ @property
272
+ def effective_sticky_mutation_authority_account_ids(self) -> frozenset[str]:
273
+ if self.sticky_mutation_authority_account_ids is None:
274
+ return frozenset(account.id for account in self.effective_continuity_owner_candidates)
275
+ return self.sticky_mutation_authority_account_ids
276
+
277
+
278
+ def _required_continuity_owner_failure(
279
+ selection_inputs: _SelectionInputs,
280
+ *,
281
+ required_account_id: str,
282
+ ) -> tuple[str, str] | None:
283
+ if selection_inputs.error_code is not None:
284
+ return None
285
+ eligible_ids = {account.id for account in selection_inputs.effective_continuity_owner_candidates} | {
286
+ account.id for account in selection_inputs.accounts
287
+ }
288
+ if required_account_id in eligible_ids:
289
+ return None
290
+ runtime_accounts = (
291
+ selection_inputs.accounts if selection_inputs.runtime_accounts is None else selection_inputs.runtime_accounts
292
+ )
293
+ if required_account_id not in {account.id for account in runtime_accounts}:
294
+ return CONTINUITY_OWNER_UNAVAILABLE, "Required continuity owner account no longer exists"
295
+ return CONTINUITY_OWNER_POLICY_CONFLICT, "Required continuity owner is outside the eligible account policy"
296
+
297
+
298
+ SelectionInputs = _SelectionInputs
299
+
300
+
301
+ class LoadBalancer:
302
+ def __init__(
303
+ self,
304
+ repo_factory: ProxyRepoFactory,
305
+ *,
306
+ encryptor: TokenEncryptor | None = None,
307
+ clock: Clock = REAL_CLOCK,
308
+ ) -> None:
309
+ self._repo_factory = repo_factory
310
+ self._encryptor = encryptor or TokenEncryptor()
311
+ self._clock = clock
312
+ self._runtime: dict[str, RuntimeState] = {}
313
+ self._runtime_lock = asyncio.Lock()
314
+ self._account_locks: dict[str, asyncio.Lock] = {}
315
+ self._account_locks_registry_lock = asyncio.Lock()
316
+ self._selection_inputs_cache = get_account_selection_cache()
317
+ # C2-2 routing/overload: the most recent request-path snapshot; paths
318
+ # without one (stream error funnel, unkeyed bridge reacquires) reuse it.
319
+ self._routing_tunables: RoutingTunables | None = None
320
+
321
+ def _resolve_routing_tunables(self, routing_tunables: RoutingTunables | None) -> RoutingTunables:
322
+ if routing_tunables is not None:
323
+ self._routing_tunables = routing_tunables
324
+ return routing_tunables
325
+ return self.current_routing_tunables()
326
+
327
+ def current_routing_tunables(self) -> RoutingTunables:
328
+ """Routing/overload knobs as of the most recent request-path snapshot."""
329
+ return self._routing_tunables or effective_routing_tunables()
330
+
331
+ async def release_account_lease(self, lease: AccountLease | None) -> None:
332
+ if lease is None:
333
+ return
334
+ async with self._runtime_lock:
335
+ self._release_account_lease_locked(lease, reason="explicit")
336
+
337
+ async def acquire_account_lease(
338
+ self,
339
+ account_id: str,
340
+ *,
341
+ kind: AccountLeaseKind,
342
+ estimated_tokens: float = 0.0,
343
+ concurrency_caps: AccountConcurrencyCaps | None = None,
344
+ api_key_id: str | None = None,
345
+ api_key_stream_fair_share_threshold_pct: int = 0,
346
+ routing_tunables: RoutingTunables | None = None,
347
+ ) -> AccountLease | None:
348
+ """Acquire a lease pinned to one account, or None on a cap denial.
349
+
350
+ Keyed stream acquires join the same per-key accounting and
351
+ congestion-gated fair-share admission as selection. A direct acquire
352
+ is pinned to one account (e.g. a warm HTTP bridge session reacquiring
353
+ its slot between turns), so the key's usable pool -- and therefore the
354
+ fair-share candidate set -- is exactly that account. Fair-share
355
+ denials raise ``ApiKeyFairShareDenialError`` so callers can
356
+ distinguish them from the plain cap denial ``None``.
357
+ """
358
+ caps = concurrency_caps or effective_account_concurrency_caps()
359
+ tunables = self._resolve_routing_tunables(routing_tunables)
360
+ async with self._runtime_lock:
361
+ self._reclaim_stale_account_leases_locked(routing_tunables=tunables)
362
+ runtime = self._runtime.setdefault(account_id, RuntimeState())
363
+ if kind == "response_create":
364
+ cap = caps.response_create_limit
365
+ if cap > 0 and runtime.inflight_response_creates >= cap:
366
+ _record_account_cap_rejection("response_create")
367
+ return None
368
+ else:
369
+ cap = caps.stream_limit
370
+ if cap > 0 and runtime.inflight_streams >= cap:
371
+ _record_account_cap_rejection("stream")
372
+ return None
373
+ denial = self._api_key_stream_fair_share_denial_locked(
374
+ api_key_id=api_key_id,
375
+ lease_kind=kind,
376
+ candidate_account_ids=(account_id,),
377
+ caps=caps,
378
+ stream_reserve_slots=0,
379
+ threshold_pct=api_key_stream_fair_share_threshold_pct,
380
+ )
381
+ if denial is not None:
382
+ raise ApiKeyFairShareDenialError(denial)
383
+ return self._acquire_account_lease_locked(
384
+ account_id,
385
+ kind=kind,
386
+ estimated_tokens=estimated_tokens,
387
+ api_key_id=api_key_id,
388
+ )
389
+
390
+ async def account_pressure_snapshot(self, account_id: str) -> tuple[int, int, float]:
391
+ async with self._runtime_lock:
392
+ runtime = self._runtime.get(account_id)
393
+ if runtime is None:
394
+ return 0, 0, 0.0
395
+ return runtime.inflight_response_creates, runtime.inflight_streams, runtime.leased_tokens
396
+
397
+ def _acquire_account_lease_locked(
398
+ self,
399
+ account_id: str,
400
+ *,
401
+ kind: AccountLeaseKind,
402
+ estimated_tokens: float,
403
+ record_selection: bool = True,
404
+ api_key_id: str | None = None,
405
+ ) -> AccountLease:
406
+ runtime = self._runtime.setdefault(account_id, RuntimeState())
407
+ lease = AccountLease(
408
+ lease_id=uuid4().hex,
409
+ account_id=account_id,
410
+ kind=kind,
411
+ acquired_at=self._clock.monotonic(),
412
+ estimated_tokens=max(0.0, estimated_tokens),
413
+ api_key_id=api_key_id,
414
+ )
415
+ if runtime.leases is None:
416
+ runtime.leases = {}
417
+ runtime.leases[lease.lease_id] = lease
418
+ if kind == "response_create":
419
+ runtime.inflight_response_creates += 1
420
+ else:
421
+ runtime.inflight_streams += 1
422
+ if api_key_id is not None:
423
+ if runtime.stream_key_inflight is None:
424
+ runtime.stream_key_inflight = {}
425
+ runtime.stream_key_inflight[api_key_id] = runtime.stream_key_inflight.get(api_key_id, 0) + 1
426
+ runtime.leased_tokens += lease.estimated_tokens
427
+ if record_selection:
428
+ runtime.last_selected_at = self._clock.time()
429
+ runtime.version += 1
430
+ _record_account_lease_acquired(kind)
431
+ _record_account_inflight_leases(account_id, runtime)
432
+ return lease
433
+
434
+ def _account_lease_allowed_locked(
435
+ self,
436
+ account_id: str,
437
+ *,
438
+ kind: AccountLeaseKind,
439
+ caps: AccountConcurrencyCaps,
440
+ stream_reserve_slots: int = 0,
441
+ ) -> bool:
442
+ runtime = self._runtime.setdefault(account_id, RuntimeState())
443
+ if kind == "response_create":
444
+ cap = caps.response_create_limit
445
+ return cap <= 0 or runtime.inflight_response_creates < cap
446
+ cap = caps.stream_limit
447
+ effective_cap = max(1, cap - max(0, stream_reserve_slots))
448
+ return cap <= 0 or runtime.inflight_streams < effective_cap
449
+
450
+ def _api_key_stream_fair_share_denial_locked(
451
+ self,
452
+ *,
453
+ api_key_id: str | None,
454
+ lease_kind: AccountLeaseKind | None,
455
+ candidate_account_ids: Collection[str],
456
+ caps: AccountConcurrencyCaps,
457
+ stream_reserve_slots: int,
458
+ threshold_pct: int,
459
+ redact_sensitive_details: bool = False,
460
+ ) -> FairShareDecision | None:
461
+ """Return the denial decision, or None when admitted or inactive.
462
+
463
+ Must be called under ``_runtime_lock``: it reads the same runtime
464
+ lease counters the admission path mutates.
465
+ """
466
+ if lease_kind != "stream" or api_key_id is None or threshold_pct <= 0:
467
+ return None
468
+ pool_capacity = effective_stream_pool_capacity(
469
+ candidate_account_count=len(candidate_account_ids),
470
+ stream_limit=caps.stream_limit,
471
+ stream_reserve_slots=stream_reserve_slots,
472
+ )
473
+ if pool_capacity <= 0:
474
+ return None
475
+ pool_inflight = 0
476
+ requester_inflight = 0
477
+ other_active_keys: set[str] = set()
478
+ for account_id in candidate_account_ids:
479
+ runtime = self._runtime.get(account_id)
480
+ if runtime is None:
481
+ continue
482
+ pool_inflight += runtime.inflight_streams
483
+ if runtime.stream_key_inflight:
484
+ requester_inflight += runtime.stream_key_inflight.get(api_key_id, 0)
485
+ other_active_keys.update(runtime.stream_key_inflight)
486
+ other_active_keys.discard(api_key_id)
487
+ decision = evaluate_stream_fair_share(
488
+ pool_capacity=pool_capacity,
489
+ pool_inflight=pool_inflight,
490
+ requester_inflight=requester_inflight,
491
+ other_active_key_count=len(other_active_keys),
492
+ threshold_pct=threshold_pct,
493
+ )
494
+ _record_stream_pool_gauges(decision.pool_capacity, decision.pool_inflight)
495
+ if decision.admitted:
496
+ return None
497
+ _record_api_key_fair_share_rejection()
498
+ logger.warning(
499
+ "API key stream fair share denial api_key_id=%s key_inflight=%s fair_share=%s "
500
+ "pool_inflight=%s pool_capacity=%s active_keys=%s",
501
+ "<redacted>" if redact_sensitive_details else api_key_id,
502
+ decision.requester_inflight,
503
+ decision.fair_share,
504
+ decision.pool_inflight,
505
+ decision.pool_capacity,
506
+ decision.active_key_count,
507
+ )
508
+ return decision
509
+
510
+ def _release_account_lease_locked(
511
+ self,
512
+ lease: AccountLease,
513
+ *,
514
+ reason: str,
515
+ redact_sensitive_details: bool = False,
516
+ ) -> bool:
517
+ runtime = self._runtime.get(lease.account_id)
518
+ if runtime is None:
519
+ return False
520
+ current = apply_lease_release(runtime, lease)
521
+ if current is None:
522
+ return False
523
+ _record_account_lease_released(current.kind, reason)
524
+ _record_account_inflight_leases(current.account_id, runtime)
525
+ if reason == "stale":
526
+ _record_account_lease_stale_reclaimed(current.kind)
527
+ logger.warning(
528
+ "Reclaimed stale account lease account_id=%s kind=%s age_seconds=%.3f",
529
+ "<redacted>" if redact_sensitive_details else current.account_id,
530
+ current.kind,
531
+ self._clock.monotonic() - current.acquired_at,
532
+ )
533
+ return True
534
+
535
+ def _reclaim_stale_account_leases_locked(
536
+ self,
537
+ *,
538
+ routing_tunables: RoutingTunables,
539
+ redact_sensitive_details: bool = False,
540
+ ) -> None:
541
+ settings = with_dashboard_overrides(get_settings())
542
+ now = self._clock.monotonic()
543
+ for runtime in self._runtime.values():
544
+ if not runtime.leases:
545
+ continue
546
+ stale = [
547
+ lease
548
+ for lease in runtime.leases.values()
549
+ if now - lease.acquired_at
550
+ >= account_lease_stale_ttl_seconds(lease.kind, settings, routing_tunables=routing_tunables)
551
+ ]
552
+ for lease in stale:
553
+ self._release_account_lease_locked(
554
+ lease,
555
+ reason="stale",
556
+ redact_sensitive_details=redact_sensitive_details,
557
+ )
558
+
559
+ def _detached_runtime_snapshot(self, *, routing_tunables: RoutingTunables) -> dict[str, RuntimeState]:
560
+ """Runtime as ordinary selection would see it, for observations that must not touch it."""
561
+ settings = with_dashboard_overrides(get_settings())
562
+ return detached_runtime_snapshot(
563
+ self._runtime,
564
+ now=self._clock.monotonic(),
565
+ stale_lease_ttl_seconds=lambda kind: account_lease_stale_ttl_seconds(
566
+ kind, settings, routing_tunables=routing_tunables
567
+ ),
568
+ )
569
+
570
+ async def select_account(
571
+ self,
572
+ sticky_key: str | None = None,
573
+ *,
574
+ sticky_kind: StickySessionKind | None = None,
575
+ reallocate_sticky: bool = False,
576
+ sticky_source: _CodexSessionSource | None = None,
577
+ legacy_sticky_key: str | None = None,
578
+ legacy_continuity_source: _CodexSessionSource | None = None,
579
+ sticky_seed_key: str | None = None,
580
+ sticky_seed_kind: StickySessionKind | None = None,
581
+ spill_bare_session_on_account_cap: bool = False,
582
+ abandon_unavailable_legacy_owner: bool = False,
583
+ require_unambiguous_account: bool = False,
584
+ sticky_max_age_seconds: int | None = None,
585
+ prefer_earlier_reset_accounts: bool = False,
586
+ prefer_earlier_reset_window: ResetPreferenceWindow = "secondary",
587
+ routing_strategy: RoutingStrategy = "capacity_weighted",
588
+ relative_availability_power: float = 2.0,
589
+ relative_availability_top_k: int = 5,
590
+ model: str | None = None,
591
+ service_tier: str | None = None,
592
+ additional_limit_name: str | None = None,
593
+ account_ids: Collection[str] | None = None,
594
+ required_account_id: str | None = None,
595
+ required_account_is_ownership_constraint: bool = False,
596
+ required_continuity_owner: bool = False,
597
+ exclude_account_ids: Collection[str] | None = None,
598
+ require_security_work_authorized: bool = False,
599
+ budget_threshold_pct: float = 95.0,
600
+ secondary_budget_threshold_pct: float = 100.0,
601
+ routing_costs_by_account_id: RoutingCostsByAccount | None = None,
602
+ lease_kind: AccountLeaseKind | None = None,
603
+ estimated_lease_tokens: float = 0.0,
604
+ stream_reserve_slots: int = 0,
605
+ traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND,
606
+ concurrency_caps: AccountConcurrencyCaps | None = None,
607
+ redact_sensitive_details: bool = False,
608
+ allow_usage_exhaustion_error: bool = True,
609
+ api_key_id: str | None = None,
610
+ api_key_stream_fair_share_threshold_pct: int = 0,
611
+ routing_tunables: RoutingTunables | None = None,
612
+ dashboard_settings: object | None = None,
613
+ ) -> AccountSelection:
614
+ if (required_account_is_ownership_constraint or required_continuity_owner) and required_account_id is None:
615
+ raise ValueError("required account ownership flags require required_account_id")
616
+ # C2-2 routing/overload: resolved once, threaded through every lock section below.
617
+ tunables = self._resolve_routing_tunables(routing_tunables)
618
+
619
+ excluded_ids = set(exclude_account_ids or ())
620
+ scoped_account_ids = None if account_ids is None else set(account_ids)
621
+ hard_affinity_owner_excluded = False
622
+ owner_restricted_selection = required_account_is_ownership_constraint or required_continuity_owner
623
+ sticky_selection_may_resolve_owner = sticky_key is not None and sticky_kind == StickySessionKind.CODEX_SESSION
624
+ # C2-3 resilience toggles: resolved from the caller's dashboard snapshot
625
+ # (the same one that produced ``concurrency_caps``), never re-read here.
626
+ resilience = resolve_resilience_toggles(dashboard_settings)
627
+
628
+ async def load_unresolved_selection_inputs() -> _SelectionInputs:
629
+ selection_inputs = await self._load_selection_inputs(
630
+ model=model,
631
+ service_tier=service_tier,
632
+ additional_limit_name=additional_limit_name,
633
+ account_ids=scoped_account_ids,
634
+ )
635
+ if require_security_work_authorized:
636
+ # Ownership scope and routing availability are separate. Even
637
+ # an already-empty routing pool must have its owner candidates
638
+ # security-filtered before conversation ambiguity is decided.
639
+ security_scope_accounts = (
640
+ selection_inputs.runtime_accounts
641
+ if selection_inputs.runtime_accounts is not None
642
+ else [
643
+ *selection_inputs.effective_continuity_owner_candidates,
644
+ *selection_inputs.accounts,
645
+ ]
646
+ )
647
+ security_authorized_account_ids = frozenset(
648
+ account.id for account in security_scope_accounts if bool(account.security_work_authorized)
649
+ )
650
+ authorized_mutation_account_ids = (
651
+ selection_inputs.effective_sticky_mutation_authority_account_ids & security_authorized_account_ids
652
+ )
653
+ authorized_owner_candidates = [
654
+ account
655
+ for account in selection_inputs.effective_continuity_owner_candidates
656
+ if bool(account.security_work_authorized)
657
+ ]
658
+ authorized_accounts = [
659
+ account for account in selection_inputs.accounts if bool(account.security_work_authorized)
660
+ ]
661
+ if selection_inputs.accounts and not authorized_accounts:
662
+ return _SelectionInputs(
663
+ accounts=[],
664
+ latest_primary={},
665
+ latest_secondary={},
666
+ latest_monthly=selection_inputs.latest_monthly,
667
+ continuity_owner_candidates=authorized_owner_candidates,
668
+ sticky_mutation_authority_account_ids=authorized_mutation_account_ids,
669
+ quota_planner_settings=selection_inputs.quota_planner_settings,
670
+ runtime_accounts=selection_inputs.runtime_accounts,
671
+ error_message="No accounts marked as authorized for security work",
672
+ error_code="no_security_work_authorized_accounts",
673
+ )
674
+ selection_inputs = _SelectionInputs(
675
+ accounts=authorized_accounts,
676
+ latest_primary=selection_inputs.latest_primary,
677
+ latest_secondary=selection_inputs.latest_secondary,
678
+ latest_monthly=selection_inputs.latest_monthly,
679
+ continuity_owner_candidates=authorized_owner_candidates,
680
+ sticky_mutation_authority_account_ids=authorized_mutation_account_ids,
681
+ quota_planner_settings=selection_inputs.quota_planner_settings,
682
+ runtime_accounts=selection_inputs.runtime_accounts,
683
+ error_message=selection_inputs.error_message,
684
+ error_code=selection_inputs.error_code,
685
+ ignore_standard_quota_account_ids=selection_inputs.ignore_standard_quota_account_ids,
686
+ ignore_standard_quota_status=selection_inputs.ignore_standard_quota_status,
687
+ persist_standard_quota_status=selection_inputs.persist_standard_quota_status,
688
+ routing_policy_override=selection_inputs.routing_policy_override,
689
+ quota_admitted_catalog_omission_account_ids=(
690
+ selection_inputs.quota_admitted_catalog_omission_account_ids
691
+ ),
692
+ )
693
+ if excluded_ids and selection_inputs.accounts:
694
+ filtered_accounts = [account for account in selection_inputs.accounts if account.id not in excluded_ids]
695
+ if require_security_work_authorized and not filtered_accounts:
696
+ return _SelectionInputs(
697
+ accounts=[],
698
+ latest_primary={},
699
+ latest_secondary={},
700
+ latest_monthly=selection_inputs.latest_monthly,
701
+ continuity_owner_candidates=selection_inputs.effective_continuity_owner_candidates,
702
+ sticky_mutation_authority_account_ids=(
703
+ selection_inputs.effective_sticky_mutation_authority_account_ids
704
+ ),
705
+ quota_planner_settings=selection_inputs.quota_planner_settings,
706
+ runtime_accounts=selection_inputs.runtime_accounts,
707
+ error_message="No accounts marked as authorized for security work",
708
+ error_code="no_security_work_authorized_accounts",
709
+ )
710
+ selection_inputs = _SelectionInputs(
711
+ accounts=filtered_accounts,
712
+ latest_primary=selection_inputs.latest_primary,
713
+ latest_secondary=selection_inputs.latest_secondary,
714
+ latest_monthly=selection_inputs.latest_monthly,
715
+ continuity_owner_candidates=selection_inputs.effective_continuity_owner_candidates,
716
+ sticky_mutation_authority_account_ids=(
717
+ selection_inputs.effective_sticky_mutation_authority_account_ids
718
+ ),
719
+ quota_planner_settings=selection_inputs.quota_planner_settings,
720
+ runtime_accounts=selection_inputs.runtime_accounts,
721
+ error_message=selection_inputs.error_message,
722
+ error_code=selection_inputs.error_code,
723
+ ignore_standard_quota_account_ids=selection_inputs.ignore_standard_quota_account_ids,
724
+ ignore_standard_quota_status=selection_inputs.ignore_standard_quota_status,
725
+ persist_standard_quota_status=selection_inputs.persist_standard_quota_status,
726
+ routing_policy_override=selection_inputs.routing_policy_override,
727
+ quota_admitted_catalog_omission_account_ids=(
728
+ selection_inputs.quota_admitted_catalog_omission_account_ids
729
+ ),
730
+ )
731
+ if required_continuity_owner:
732
+ assert required_account_id is not None
733
+ failure = _required_continuity_owner_failure(
734
+ selection_inputs,
735
+ required_account_id=required_account_id,
736
+ )
737
+ if failure is not None:
738
+ error_code, error_message = failure
739
+ return replace(
740
+ selection_inputs,
741
+ accounts=[],
742
+ latest_primary={},
743
+ latest_secondary={},
744
+ error_message=error_message,
745
+ error_code=error_code,
746
+ )
747
+ return selection_inputs
748
+
749
+ async def load_selection_inputs() -> _SelectionInputs:
750
+ # Applied after exclusion/security filtering and on every reload,
751
+ # so retries and filtered pools keep the dashboard soft-drain value.
752
+ return replace(await load_unresolved_selection_inputs(), soft_drain_enabled=resilience.soft_drain_enabled)
753
+
754
+ selection_inputs = await load_selection_inputs()
755
+ caps = concurrency_caps or effective_account_concurrency_caps()
756
+ circuit_breaker_open = _is_upstream_circuit_breaker_open(resilience.circuit_breaker_enabled)
757
+ if circuit_breaker_open:
758
+ set_degraded("upstream circuit breaker is open")
759
+ elif (
760
+ not owner_restricted_selection
761
+ and not sticky_selection_may_resolve_owner
762
+ and (selection_inputs.accounts or selection_inputs.error_code is not None)
763
+ ):
764
+ set_normal()
765
+
766
+ if selection_inputs.error_code in {
767
+ CONTINUITY_OWNER_UNAVAILABLE,
768
+ CONTINUITY_OWNER_POLICY_CONFLICT,
769
+ }:
770
+ return AccountSelection(
771
+ account=None,
772
+ error_message=selection_inputs.error_message,
773
+ error_code=selection_inputs.error_code,
774
+ continuity_owner_no_longer_exists=selection_inputs.error_code == CONTINUITY_OWNER_UNAVAILABLE,
775
+ )
776
+
777
+ selected_snapshot: Account | None = None
778
+ error_message: str | None = None
779
+ selected_lease: AccountLease | None = None
780
+ selection_error_code: str | None = None
781
+ selection_resets_at: int | None = None
782
+ legacy_existing_account_id: str | None = None
783
+ legacy_abandoned_account_id: str | None = None
784
+ sticky_seed_account_id: str | None = None
785
+ initial_sticky_owner_lookup: StickyOwnerLookup | None = None
786
+ needs_owner_lookups = (
787
+ legacy_sticky_key is not None
788
+ or (sticky_seed_key is not None and sticky_seed_kind is not None)
789
+ or (sticky_key is not None and sticky_kind is not None)
790
+ )
791
+ if needs_owner_lookups:
792
+ # One shared session serves the legacy/seed/first-sticky owner
793
+ # lookups. The SELECTs stay separate on purpose so the per-source
794
+ # predicate semantics of get_account_id_and_abandonment (tombstone
795
+ # visibility, max_age handling) are untouched; the saving is the
796
+ # 2-3 extra pool checkouts + session create/teardown lifecycles
797
+ # per request. Each later source still starts a fresh read
798
+ # transaction (release_read_snapshot): on SQLite/WAL the shared
799
+ # session would otherwise pin one snapshot at the first SELECT
800
+ # and hide a hard sticky or seed owner committed concurrently
801
+ # between the reads, letting selection overwrite that mapping.
802
+ async with self._repo_factory() as repos:
803
+ owner_snapshot_pinned = False
804
+ if legacy_sticky_key is not None:
805
+ legacy_owner_lookup = await repos.sticky_sessions.get_account_id_and_abandonment(
806
+ legacy_sticky_key,
807
+ kind=StickySessionKind.CODEX_SESSION,
808
+ # Raw rows may be historical turn-state ownership. The
809
+ # bounded thread TTL must never age out that hard evidence.
810
+ max_age_seconds=None,
811
+ # Process-session raw text is session_header even when
812
+ # request locality is thread_header. Thread-only raw keys
813
+ # keep thread_header so a session_header tombstone cannot
814
+ # hide a distinct thread owner.
815
+ continuity_source=legacy_continuity_source or "session_header",
816
+ )
817
+ legacy_existing_account_id = legacy_owner_lookup.account_id
818
+ abandoned_account_id = legacy_owner_lookup.abandoned_account_id
819
+ if legacy_owner_lookup.continuity_abandoned is True and isinstance(abandoned_account_id, str):
820
+ legacy_abandoned_account_id = abandoned_account_id
821
+ if required_account_id is not None and (
822
+ legacy_existing_account_id is not None and legacy_existing_account_id != required_account_id
823
+ ):
824
+ # The required owner came from a file/response/bridge index,
825
+ # while the raw row may be legacy turn-state ownership. Neither
826
+ # source can be discarded or rewritten to resolve a conflict.
827
+ return AccountSelection(
828
+ account=None,
829
+ error_message="Account-owned continuity sources conflict; retry the logical turn",
830
+ error_code="continuity_owner_conflict",
831
+ )
832
+ owner_snapshot_pinned = True
833
+ if sticky_seed_key is not None and sticky_seed_kind is not None:
834
+ if owner_snapshot_pinned:
835
+ await repos.sticky_sessions.release_read_snapshot()
836
+ sticky_seed_account_id = await repos.sticky_sessions.get_account_id(
837
+ sticky_seed_key,
838
+ kind=sticky_seed_kind,
839
+ )
840
+ owner_snapshot_pinned = True
841
+ if sticky_key is not None and sticky_kind is not None:
842
+ if owner_snapshot_pinned:
843
+ await repos.sticky_sessions.release_read_snapshot()
844
+ # First-iteration owner read for run_sticky_selection_path,
845
+ # hoisted here so it shares this session. The selection
846
+ # loop consumes it exactly once; every retry (including
847
+ # post-reset attempts) still re-reads fresh ownership
848
+ # evidence through its own repo bundle.
849
+ initial_sticky_owner_lookup = await repos.sticky_sessions.get_account_id_and_abandonment(
850
+ sticky_key,
851
+ kind=sticky_kind,
852
+ max_age_seconds=sticky_max_age_seconds,
853
+ continuity_source=sticky_source,
854
+ )
855
+ # Resolve uniqueness from the model/API-key/security-scoped pool before
856
+ # runtime health, budget, or cap filtering. Transient pressure cannot
857
+ # prove that another candidate does not own an upstream conversation.
858
+ if (
859
+ require_unambiguous_account
860
+ and sticky_key is None
861
+ and legacy_existing_account_id is None
862
+ and len(selection_inputs.effective_continuity_owner_candidates) != 1
863
+ ):
864
+ return AccountSelection(
865
+ account=None,
866
+ error_message=_AMBIGUOUS_CONVERSATION_OWNER_MESSAGE,
867
+ error_code=_AMBIGUOUS_CONVERSATION_OWNER_CODE,
868
+ )
869
+ # Transient routing errors are secondary to ownership ambiguity. An
870
+ # empty additional-quota pool cannot prove which account owns a
871
+ # conversation that was ambiguous before that filter ran.
872
+ if selection_inputs.error_code is not None and not selection_inputs.accounts:
873
+ return AccountSelection(
874
+ account=None,
875
+ error_message=selection_inputs.error_message,
876
+ error_code=selection_inputs.error_code,
877
+ )
878
+ if sticky_key is None:
879
+ unbound_outcome = await run_unbound_selection_path(
880
+ self,
881
+ request=UnboundSelectionRequest(
882
+ prefer_earlier_reset_accounts=prefer_earlier_reset_accounts,
883
+ prefer_earlier_reset_window=prefer_earlier_reset_window,
884
+ routing_strategy=routing_strategy,
885
+ relative_availability_power=relative_availability_power,
886
+ relative_availability_top_k=relative_availability_top_k,
887
+ required_account_id=required_account_id,
888
+ budget_threshold_pct=budget_threshold_pct,
889
+ secondary_budget_threshold_pct=secondary_budget_threshold_pct,
890
+ routing_costs_by_account_id=routing_costs_by_account_id,
891
+ lease_kind=lease_kind,
892
+ estimated_lease_tokens=estimated_lease_tokens,
893
+ stream_reserve_slots=stream_reserve_slots,
894
+ traffic_class=traffic_class,
895
+ concurrency_caps=caps,
896
+ redact_sensitive_details=redact_sensitive_details,
897
+ routing_tunables=tunables,
898
+ api_key_id=api_key_id,
899
+ api_key_stream_fair_share_threshold_pct=api_key_stream_fair_share_threshold_pct,
900
+ model=model,
901
+ selection_inputs=selection_inputs,
902
+ reload_inputs=load_selection_inputs,
903
+ record_account_cap_rejection=_record_account_cap_rejection,
904
+ allow_usage_exhaustion_error=allow_usage_exhaustion_error,
905
+ ),
906
+ )
907
+ selection_inputs = unbound_outcome.selection_inputs
908
+ selected_snapshot = unbound_outcome.selected_snapshot
909
+ selected_lease = unbound_outcome.selected_lease
910
+ error_message = unbound_outcome.error_message
911
+ selection_error_code = unbound_outcome.error_code
912
+ selection_resets_at = unbound_outcome.resets_at
913
+ if unbound_outcome.disposition == "direct_error":
914
+ return AccountSelection(
915
+ account=None,
916
+ error_message=error_message,
917
+ error_code=selection_error_code,
918
+ )
919
+ if (
920
+ selected_snapshot is not None
921
+ and selected_lease is not None
922
+ and sticky_seed_key is not None
923
+ and sticky_seed_kind is not None
924
+ and sticky_seed_account_id is None
925
+ ):
926
+ # Required-owner selection bypasses the thread row, but a
927
+ # first-ever process preference still has to land so later
928
+ # unpinned siblings inherit that exact owner.
929
+ try:
930
+ async with self._repo_factory() as repos:
931
+ await repos.sticky_sessions.insert_if_absent(
932
+ sticky_seed_key,
933
+ selected_snapshot.id,
934
+ sticky_seed_kind,
935
+ )
936
+ except BaseException:
937
+ await self.release_account_lease(selected_lease)
938
+ selected_lease = None
939
+ raise
940
+ else:
941
+ sticky_outcome = await run_sticky_selection_path(
942
+ self,
943
+ request=StickySelectionRequest(
944
+ routing_tunables=tunables,
945
+ sticky_key=sticky_key,
946
+ sticky_kind=sticky_kind,
947
+ reallocate_sticky=reallocate_sticky,
948
+ sticky_source=sticky_source,
949
+ legacy_sticky_key=legacy_sticky_key,
950
+ legacy_existing_account_id=legacy_existing_account_id,
951
+ legacy_abandoned_account_id=legacy_abandoned_account_id,
952
+ sticky_seed_key=sticky_seed_key,
953
+ sticky_seed_kind=sticky_seed_kind,
954
+ sticky_seed_account_id=sticky_seed_account_id,
955
+ spill_bare_session_on_account_cap=spill_bare_session_on_account_cap,
956
+ abandon_unavailable_legacy_owner=abandon_unavailable_legacy_owner,
957
+ require_unambiguous_account=require_unambiguous_account,
958
+ sticky_max_age_seconds=sticky_max_age_seconds,
959
+ prefer_earlier_reset_accounts=prefer_earlier_reset_accounts,
960
+ prefer_earlier_reset_window=prefer_earlier_reset_window,
961
+ routing_strategy=routing_strategy,
962
+ relative_availability_power=relative_availability_power,
963
+ relative_availability_top_k=relative_availability_top_k,
964
+ required_account_id=required_account_id,
965
+ budget_threshold_pct=budget_threshold_pct,
966
+ secondary_budget_threshold_pct=secondary_budget_threshold_pct,
967
+ routing_costs_by_account_id=routing_costs_by_account_id,
968
+ lease_kind=lease_kind,
969
+ estimated_lease_tokens=estimated_lease_tokens,
970
+ stream_reserve_slots=stream_reserve_slots,
971
+ traffic_class=traffic_class,
972
+ concurrency_caps=caps,
973
+ redact_sensitive_details=redact_sensitive_details,
974
+ api_key_id=api_key_id,
975
+ api_key_stream_fair_share_threshold_pct=api_key_stream_fair_share_threshold_pct,
976
+ exclude_account_ids=frozenset(excluded_ids),
977
+ model=model,
978
+ selection_inputs=selection_inputs,
979
+ reload_inputs=load_selection_inputs,
980
+ record_account_cap_rejection=_record_account_cap_rejection,
981
+ allow_usage_exhaustion_error=allow_usage_exhaustion_error,
982
+ initial_sticky_owner_lookup=initial_sticky_owner_lookup,
983
+ ),
984
+ )
985
+ selection_inputs = sticky_outcome.selection_inputs
986
+ selected_snapshot = sticky_outcome.selected_snapshot
987
+ selected_lease = sticky_outcome.selected_lease
988
+ error_message = sticky_outcome.error_message
989
+ selection_error_code = sticky_outcome.error_code
990
+ selection_resets_at = sticky_outcome.resets_at
991
+ hard_affinity_owner_excluded = sticky_outcome.hard_affinity_owner_excluded
992
+ if sticky_outcome.disposition == "direct_error":
993
+ return AccountSelection(
994
+ account=None,
995
+ error_message=error_message,
996
+ error_code=selection_error_code,
997
+ )
998
+
999
+ if selected_snapshot is None:
1000
+ logger.warning(
1001
+ "No account selected strategy=%s sticky=%s model=%s error=%s",
1002
+ routing_strategy,
1003
+ bool(sticky_key),
1004
+ model,
1005
+ error_message,
1006
+ )
1007
+
1008
+ if selected_snapshot is None:
1009
+ owner_restricted_selection = owner_restricted_selection or selection_error_code == "hard_affinity_saturated"
1010
+ opportunistic_policy_blocked = (
1011
+ traffic_class == TRAFFIC_CLASS_OPPORTUNISTIC
1012
+ and error_message is not None
1013
+ and error_message.startswith("opportunistic burn window closed")
1014
+ )
1015
+ if opportunistic_policy_blocked:
1016
+ return AccountSelection(
1017
+ account=None,
1018
+ error_message=error_message,
1019
+ error_code=OPPORTUNISTIC_BURN_WINDOW_CLOSED,
1020
+ )
1021
+ if required_continuity_owner and selection_error_code is None:
1022
+ selection_error_code = CONTINUITY_OWNER_UNAVAILABLE
1023
+ if traffic_class == TRAFFIC_CLASS_OPPORTUNISTIC and error_message and selection_error_code is None:
1024
+ return AccountSelection(
1025
+ account=None,
1026
+ error_message=error_message,
1027
+ error_code=OPPORTUNISTIC_BURN_WINDOW_CLOSED,
1028
+ )
1029
+ if error_message == "No available accounts" and not owner_restricted_selection:
1030
+ set_degraded("all upstream accounts are unavailable")
1031
+ error_message = _format_degraded_error_message(error_message)
1032
+ elif (
1033
+ not owner_restricted_selection
1034
+ and not circuit_breaker_open
1035
+ and (selection_inputs.accounts or selection_inputs.error_code is not None)
1036
+ ):
1037
+ set_normal()
1038
+ return AccountSelection(
1039
+ account=None,
1040
+ error_message=error_message,
1041
+ error_code=selection_error_code,
1042
+ resets_at=selection_resets_at,
1043
+ hard_affinity_owner_excluded=hard_affinity_owner_excluded,
1044
+ )
1045
+ if not circuit_breaker_open:
1046
+ set_normal()
1047
+ logger.info(
1048
+ "Selected account_id=%s strategy=%s sticky=%s model=%s",
1049
+ "<redacted>" if redact_sensitive_details else selected_snapshot.id,
1050
+ routing_strategy,
1051
+ bool(sticky_key),
1052
+ model,
1053
+ )
1054
+ return AccountSelection(
1055
+ account=selected_snapshot,
1056
+ error_message=None,
1057
+ error_code=None,
1058
+ lease=selected_lease,
1059
+ catalog_omission_quota_admission=_catalog_omission_quota_admission(
1060
+ account_id=selected_snapshot.id,
1061
+ model=model,
1062
+ service_tier=service_tier,
1063
+ additional_limit_name=additional_limit_name,
1064
+ quota_admitted_catalog_omission_account_ids=(
1065
+ selection_inputs.quota_admitted_catalog_omission_account_ids
1066
+ ),
1067
+ ),
1068
+ )
1069
+
1070
+ def _reserve_due_probe_locked(
1071
+ self,
1072
+ states: list[AccountState],
1073
+ *,
1074
+ prefer_earlier_reset: bool,
1075
+ prefer_earlier_reset_window: ResetPreferenceWindow,
1076
+ routing_strategy: RoutingStrategy,
1077
+ relative_availability_power: float,
1078
+ relative_availability_top_k: int,
1079
+ traffic_class: TrafficClass,
1080
+ routing_costs_by_account_id: RoutingCostsByAccount | None,
1081
+ ) -> ProbeReservation | None:
1082
+ if routing_strategy in ("sequential_drain", "reset_drain", "single_account"):
1083
+ return None
1084
+ result = select_account(
1085
+ states,
1086
+ now=self._clock.time(),
1087
+ prefer_earlier_reset=prefer_earlier_reset,
1088
+ prefer_earlier_reset_window=prefer_earlier_reset_window,
1089
+ routing_strategy=routing_strategy,
1090
+ recovery_probe_only=True,
1091
+ relative_availability_power=relative_availability_power,
1092
+ relative_availability_top_k=relative_availability_top_k,
1093
+ traffic_class=traffic_class,
1094
+ routing_costs=routing_costs_by_account_id,
1095
+ )
1096
+ if result.account is None:
1097
+ return None
1098
+ runtime = self._runtime.get(result.account.account_id)
1099
+ if runtime is None:
1100
+ return None
1101
+ if runtime.health_tier != result.account.health_tier:
1102
+ return None
1103
+ if runtime.last_selected_at != result.account.last_selected_at:
1104
+ return None
1105
+ # Keep the current state snapshot due for this request while making a
1106
+ # concurrent snapshot observe the reservation before sticky DB I/O.
1107
+ # This is not a health observation, so it must not advance ``version``
1108
+ # and invalidate an operator Force Probe that is loading usage.
1109
+ previous_last_selected_at = runtime.last_selected_at
1110
+ reserved_at = self._clock.time()
1111
+ runtime.last_selected_at = reserved_at
1112
+ return ProbeReservation(
1113
+ account_id=result.account.account_id,
1114
+ previous_last_selected_at=previous_last_selected_at,
1115
+ reserved_at=reserved_at,
1116
+ expected_runtime_version=runtime.version,
1117
+ )
1118
+
1119
+ def _probe_reservation_current_locked(self, reservation: ProbeReservation | None) -> bool:
1120
+ if reservation is None:
1121
+ return False
1122
+ runtime = self._runtime.get(reservation.account_id)
1123
+ return bool(
1124
+ runtime is not None
1125
+ and runtime.last_selected_at == reservation.reserved_at
1126
+ and runtime.version == reservation.expected_runtime_version
1127
+ )
1128
+
1129
+ def _release_due_probe_reservation_locked(self, reservation: ProbeReservation | None) -> None:
1130
+ if reservation is None:
1131
+ return
1132
+ runtime = self._runtime.get(reservation.account_id)
1133
+ # An actual concurrent selection replaces this exact timestamp and
1134
+ # owns the admission. Unrelated runtime changes may advance version but
1135
+ # must not turn a temporary reservation into a consumed probe interval.
1136
+ if runtime is None or runtime.last_selected_at != reservation.reserved_at:
1137
+ return
1138
+ runtime.last_selected_at = reservation.previous_last_selected_at
1139
+
1140
+ def _commit_due_probe_reservation_locked(self, reservation: ProbeReservation | None) -> bool:
1141
+ if reservation is None:
1142
+ return False
1143
+ runtime = self._runtime.get(reservation.account_id)
1144
+ if runtime is None or not self._probe_reservation_current_locked(reservation):
1145
+ return False
1146
+ # Only a selection that survived sticky persistence and final local
1147
+ # admission consumes the quiet interval. Unlike reserve/release, this
1148
+ # committed observation must invalidate older Force Probe settlement.
1149
+ runtime.last_selected_at = self._clock.time()
1150
+ runtime.version += 1
1151
+ runtime.health_version += 1
1152
+ return True
1153
+
1154
+ def _sync_committed_probe_state_locked(
1155
+ self,
1156
+ reservation: ProbeReservation,
1157
+ account_map: Mapping[str, Account],
1158
+ states: Collection[AccountState],
1159
+ ) -> None:
1160
+ account = account_map.get(reservation.account_id)
1161
+ if account is None:
1162
+ return
1163
+ for state in states:
1164
+ if state.account_id == reservation.account_id:
1165
+ self._sync_runtime_state(account, state)
1166
+ return
1167
+
1168
+ async def _load_selection_inputs(
1169
+ self,
1170
+ *,
1171
+ model: str | None,
1172
+ service_tier: str | None = None,
1173
+ additional_limit_name: str | None = None,
1174
+ account_ids: Collection[str] | None = None,
1175
+ ) -> _SelectionInputs:
1176
+ mapped_limit_name = _gated_limit_name_for_model(model)
1177
+ effective_limit_name = additional_limit_name or mapped_limit_name
1178
+ additional_quota_routing_policies: dict[str, str] = {}
1179
+ if effective_limit_name is not None:
1180
+ additional_quota_routing_policies = await _load_dashboard_additional_quota_routing_overrides()
1181
+ additional_quota_routing_policies_cache_key = json.dumps(
1182
+ additional_quota_routing_policies,
1183
+ sort_keys=True,
1184
+ separators=(",", ":"),
1185
+ )
1186
+ cache_key = (
1187
+ model,
1188
+ service_tier,
1189
+ additional_limit_name,
1190
+ additional_quota_routing_policies_cache_key,
1191
+ None if account_ids is None else tuple(sorted(set(account_ids))),
1192
+ )
1193
+ cached = await self._selection_inputs_cache.get(cache_key)
1194
+ if cached is not None:
1195
+ return _clone_selection_inputs(cached)
1196
+
1197
+ load_generation = self._selection_inputs_cache.generation
1198
+
1199
+ async with self._repo_factory() as repos:
1200
+ all_accounts = await repos.accounts.list_accounts()
1201
+ quota_planner_repo = getattr(repos, "quota_planner", None)
1202
+ get_quota_planner_settings = getattr(quota_planner_repo, "get_settings", None)
1203
+ if callable(get_quota_planner_settings):
1204
+ try:
1205
+ settings_result = get_quota_planner_settings()
1206
+ quota_planner_settings = (
1207
+ await settings_result if inspect.isawaitable(settings_result) else settings_result
1208
+ )
1209
+ if not isinstance(quota_planner_settings, PlannerSettings):
1210
+ quota_planner_settings = PlannerSettings()
1211
+ except Exception:
1212
+ logger.warning("Failed to load quota planner settings; using defaults", exc_info=True)
1213
+ quota_planner_settings = PlannerSettings()
1214
+ else:
1215
+ quota_planner_settings = PlannerSettings()
1216
+ ignore_standard_quota_status = effective_limit_name is not None
1217
+ routing_policy_override = _additional_quota_routing_policy_override(
1218
+ effective_limit_name,
1219
+ additional_quota_routing_policies,
1220
+ )
1221
+ scoped_accounts = all_accounts
1222
+ if account_ids is not None:
1223
+ allowed_account_ids = set(account_ids)
1224
+ scoped_accounts = [account for account in scoped_accounts if account.id in allowed_account_ids]
1225
+ sticky_mutation_authority_account_ids = frozenset(account.id for account in scoped_accounts)
1226
+ accounts = _selectable_accounts(scoped_accounts)
1227
+ pre_model_filter_accounts = accounts
1228
+ model_catalog_omitted_account_ids: frozenset[str] = frozenset()
1229
+ applied_service_tier: str | None = None
1230
+ if model and _mapped_model_has_registry_entry(model):
1231
+ continuity_owner_candidates = _filter_accounts_for_model(
1232
+ scoped_accounts,
1233
+ model,
1234
+ service_tier=service_tier,
1235
+ )
1236
+ canonical_quota_can_override_account_catalog = (
1237
+ additional_limit_name is None and mapped_limit_name is not None
1238
+ )
1239
+ model_filter = _filter_accounts_for_model_with_catalog_evidence(
1240
+ pre_model_filter_accounts,
1241
+ model,
1242
+ service_tier=service_tier,
1243
+ additional_quota_can_override_account_catalog=canonical_quota_can_override_account_catalog,
1244
+ )
1245
+ accounts = model_filter.accounts
1246
+ general_model_account_ids = model_filter.general_model_account_ids
1247
+ applied_service_tier = model_filter.applied_service_tier
1248
+ if canonical_quota_can_override_account_catalog and general_model_account_ids is not None:
1249
+ model_catalog_omitted_account_ids = frozenset(
1250
+ account.id for account in accounts if account.id not in general_model_account_ids
1251
+ )
1252
+ else:
1253
+ # Administrative/runtime status affects routability, not who
1254
+ # may own account-scoped upstream state. Capture this pool
1255
+ # before PAUSED/DEACTIVATED/etc. can manufacture uniqueness.
1256
+ continuity_owner_candidates = scoped_accounts
1257
+ if model and not accounts:
1258
+ if not all_accounts:
1259
+ selection_inputs = _SelectionInputs(
1260
+ accounts=[],
1261
+ latest_primary={},
1262
+ latest_secondary={},
1263
+ latest_monthly={},
1264
+ continuity_owner_candidates=[
1265
+ _clone_account(account) for account in continuity_owner_candidates
1266
+ ],
1267
+ sticky_mutation_authority_account_ids=sticky_mutation_authority_account_ids,
1268
+ quota_planner_settings=quota_planner_settings,
1269
+ runtime_accounts=[_clone_account(account) for account in all_accounts],
1270
+ )
1271
+ await self._selection_inputs_cache.set(
1272
+ _clone_selection_inputs(selection_inputs), key=cache_key, generation=load_generation
1273
+ )
1274
+ return selection_inputs
1275
+ if not pre_model_filter_accounts:
1276
+ selection_inputs = _SelectionInputs(
1277
+ accounts=[],
1278
+ latest_primary={},
1279
+ latest_secondary={},
1280
+ latest_monthly={},
1281
+ continuity_owner_candidates=[],
1282
+ sticky_mutation_authority_account_ids=sticky_mutation_authority_account_ids,
1283
+ quota_planner_settings=quota_planner_settings,
1284
+ runtime_accounts=[_clone_account(account) for account in all_accounts],
1285
+ )
1286
+ await self._selection_inputs_cache.set(
1287
+ _clone_selection_inputs(selection_inputs), key=cache_key, generation=load_generation
1288
+ )
1289
+ return selection_inputs
1290
+ if continuity_owner_candidates:
1291
+ selection_inputs = _SelectionInputs(
1292
+ accounts=[],
1293
+ latest_primary={},
1294
+ latest_secondary={},
1295
+ latest_monthly={},
1296
+ continuity_owner_candidates=[
1297
+ _clone_account(account) for account in continuity_owner_candidates
1298
+ ],
1299
+ sticky_mutation_authority_account_ids=sticky_mutation_authority_account_ids,
1300
+ quota_planner_settings=quota_planner_settings,
1301
+ runtime_accounts=[_clone_account(account) for account in all_accounts],
1302
+ )
1303
+ await self._selection_inputs_cache.set(
1304
+ _clone_selection_inputs(selection_inputs), key=cache_key, generation=load_generation
1305
+ )
1306
+ return selection_inputs
1307
+ selection_inputs = _SelectionInputs(
1308
+ accounts=[],
1309
+ latest_primary={},
1310
+ latest_secondary={},
1311
+ latest_monthly={},
1312
+ continuity_owner_candidates=[_clone_account(account) for account in continuity_owner_candidates],
1313
+ sticky_mutation_authority_account_ids=sticky_mutation_authority_account_ids,
1314
+ quota_planner_settings=quota_planner_settings,
1315
+ runtime_accounts=[_clone_account(account) for account in all_accounts],
1316
+ error_message=(
1317
+ f"No accounts with a plan supporting model '{model}' at service tier '{applied_service_tier}'"
1318
+ if applied_service_tier is not None
1319
+ else f"No accounts with a plan supporting model '{model}'"
1320
+ ),
1321
+ error_code=NO_PLAN_SUPPORT_FOR_MODEL,
1322
+ )
1323
+ await self._selection_inputs_cache.set(
1324
+ _clone_selection_inputs(selection_inputs), key=cache_key, generation=load_generation
1325
+ )
1326
+ return selection_inputs
1327
+
1328
+ if effective_limit_name:
1329
+ additional_quota_candidates = accounts
1330
+ additional_filter = await self._filter_accounts_for_additional_limit(
1331
+ additional_quota_candidates,
1332
+ model=model,
1333
+ limit_name=effective_limit_name,
1334
+ explicit_limit=additional_limit_name is not None,
1335
+ repos=repos,
1336
+ require_fresh_evidence_account_ids=model_catalog_omitted_account_ids,
1337
+ )
1338
+ accounts = additional_filter.accounts
1339
+ if not accounts and all_accounts_require_reauthentication(
1340
+ additional_quota_candidates,
1341
+ self._encryptor,
1342
+ now=self._clock.time(),
1343
+ ):
1344
+ accounts = additional_quota_candidates
1345
+ elif not accounts:
1346
+ selection_inputs = _SelectionInputs(
1347
+ accounts=[],
1348
+ latest_primary={},
1349
+ latest_secondary={},
1350
+ latest_monthly={},
1351
+ continuity_owner_candidates=[
1352
+ _clone_account(account) for account in continuity_owner_candidates
1353
+ ],
1354
+ sticky_mutation_authority_account_ids=sticky_mutation_authority_account_ids,
1355
+ quota_planner_settings=quota_planner_settings,
1356
+ runtime_accounts=[_clone_account(account) for account in all_accounts],
1357
+ error_message=additional_filter.error_message,
1358
+ error_code=additional_filter.error_code,
1359
+ )
1360
+ await self._selection_inputs_cache.set(
1361
+ _clone_selection_inputs(selection_inputs), key=cache_key, generation=load_generation
1362
+ )
1363
+ return selection_inputs
1364
+ if not accounts:
1365
+ selection_inputs = _SelectionInputs(
1366
+ accounts=[],
1367
+ latest_primary={},
1368
+ latest_secondary={},
1369
+ latest_monthly={},
1370
+ continuity_owner_candidates=[_clone_account(account) for account in continuity_owner_candidates],
1371
+ sticky_mutation_authority_account_ids=sticky_mutation_authority_account_ids,
1372
+ quota_planner_settings=quota_planner_settings,
1373
+ runtime_accounts=[_clone_account(account) for account in all_accounts],
1374
+ )
1375
+ await self._selection_inputs_cache.set(
1376
+ _clone_selection_inputs(selection_inputs), key=cache_key, generation=load_generation
1377
+ )
1378
+ return selection_inputs
1379
+
1380
+ # These share one AsyncSession: concurrent execution on a single
1381
+ # session is unsafe (asyncpg) and gains nothing — the driver
1382
+ # serializes statements per connection anyway.
1383
+ standard_latest_primary = await repos.usage.latest_by_account()
1384
+ standard_latest_secondary = await repos.usage.latest_by_account(window="secondary")
1385
+ latest_monthly = await repos.usage.latest_by_account(window="monthly")
1386
+ if effective_limit_name:
1387
+ model_allowed_plans = get_model_registry().plan_types_for_model(model) if model else None
1388
+ latest_primary = additional_filter.latest_primary
1389
+ latest_secondary = additional_filter.latest_secondary
1390
+ quota_scoped_account_ids = frozenset(
1391
+ account.id
1392
+ for account in accounts
1393
+ if additional_limit_name is not None
1394
+ or (
1395
+ model_allowed_plans is not None
1396
+ and normalize_account_plan_type(account.plan_type) not in _ADDITIONAL_QUOTA_EXEMPT_PLAN_TYPES
1397
+ and account_plan_matches_allowed(
1398
+ account.plan_type,
1399
+ model_allowed_plans,
1400
+ )
1401
+ )
1402
+ )
1403
+ latest_primary: dict[str, UsageHistory | AdditionalUsageHistory] = dict(standard_latest_primary)
1404
+ latest_secondary: dict[str, UsageHistory | AdditionalUsageHistory] = dict(standard_latest_secondary)
1405
+ for account_id in quota_scoped_account_ids:
1406
+ latest_primary.pop(account_id, None)
1407
+ latest_secondary.pop(account_id, None)
1408
+ if account_id in additional_filter.latest_primary:
1409
+ latest_primary[account_id] = additional_filter.latest_primary[account_id]
1410
+ if account_id in additional_filter.latest_secondary:
1411
+ latest_secondary[account_id] = additional_filter.latest_secondary[account_id]
1412
+ ignore_standard_quota_account_ids = quota_scoped_account_ids
1413
+ else:
1414
+ latest_primary = standard_latest_primary
1415
+ latest_secondary = standard_latest_secondary
1416
+ ignore_standard_quota_account_ids = frozenset()
1417
+ quota_admitted_catalog_omission_account_ids = frozenset(
1418
+ account.id for account in accounts if account.id in model_catalog_omitted_account_ids
1419
+ )
1420
+ selection_inputs = _SelectionInputs(
1421
+ accounts=[_clone_account(account) for account in accounts],
1422
+ latest_primary={account_id: clone_row(entry) for account_id, entry in latest_primary.items()},
1423
+ latest_secondary={account_id: clone_row(entry) for account_id, entry in latest_secondary.items()},
1424
+ latest_monthly={account_id: clone_row(entry) for account_id, entry in latest_monthly.items()},
1425
+ continuity_owner_candidates=[_clone_account(account) for account in continuity_owner_candidates],
1426
+ sticky_mutation_authority_account_ids=sticky_mutation_authority_account_ids,
1427
+ quota_planner_settings=quota_planner_settings,
1428
+ runtime_accounts=[_clone_account(account) for account in all_accounts],
1429
+ ignore_standard_quota_account_ids=ignore_standard_quota_account_ids,
1430
+ ignore_standard_quota_status=ignore_standard_quota_status,
1431
+ persist_standard_quota_status=True,
1432
+ routing_policy_override=routing_policy_override,
1433
+ quota_admitted_catalog_omission_account_ids=quota_admitted_catalog_omission_account_ids,
1434
+ )
1435
+ await self._selection_inputs_cache.set(
1436
+ _clone_selection_inputs(selection_inputs), key=cache_key, generation=load_generation
1437
+ )
1438
+ return selection_inputs
1439
+
1440
+ async def check_opportunistic_admission(
1441
+ self,
1442
+ *,
1443
+ model: str | None,
1444
+ account_ids: Collection[str] | None,
1445
+ prefer_earlier_reset_accounts: bool,
1446
+ routing_strategy: RoutingStrategy,
1447
+ budget_threshold_pct: float,
1448
+ prefer_earlier_reset_window: ResetPreferenceWindow = "secondary",
1449
+ secondary_budget_threshold_pct: float = 100.0,
1450
+ lease_kind: AccountLeaseKind | None = None,
1451
+ concurrency_caps: AccountConcurrencyCaps | None = None,
1452
+ stream_reserve_slots: int = 0,
1453
+ service_tier: str | None = None,
1454
+ observe_only: bool = False,
1455
+ routing_tunables: RoutingTunables | None = None,
1456
+ dashboard_settings: object | None = None,
1457
+ ) -> AccountSelection:
1458
+ outcome = await run_opportunistic_admission(
1459
+ self,
1460
+ request=OpportunisticAdmissionRequest(
1461
+ model=model,
1462
+ service_tier=service_tier,
1463
+ account_ids=account_ids,
1464
+ prefer_earlier_reset_accounts=prefer_earlier_reset_accounts,
1465
+ prefer_earlier_reset_window=prefer_earlier_reset_window,
1466
+ routing_strategy=routing_strategy,
1467
+ budget_threshold_pct=budget_threshold_pct,
1468
+ secondary_budget_threshold_pct=secondary_budget_threshold_pct,
1469
+ lease_kind=lease_kind,
1470
+ concurrency_caps=concurrency_caps or effective_account_concurrency_caps(),
1471
+ stream_reserve_slots=stream_reserve_slots,
1472
+ record_account_cap_rejection=_record_account_cap_rejection,
1473
+ build_states=_build_states,
1474
+ routing_tunables=self._resolve_routing_tunables(routing_tunables),
1475
+ observe_only=observe_only,
1476
+ # C2-3 resilience toggles: from the caller's dashboard snapshot.
1477
+ soft_drain_enabled=resolve_resilience_toggles(dashboard_settings).soft_drain_enabled,
1478
+ ),
1479
+ )
1480
+ return AccountSelection(
1481
+ account=outcome.account,
1482
+ error_message=outcome.error_message,
1483
+ error_code=outcome.error_code,
1484
+ resets_at=outcome.resets_at,
1485
+ )
1486
+
1487
+ async def _filter_accounts_for_additional_limit(
1488
+ self,
1489
+ accounts: list[Account],
1490
+ *,
1491
+ model: str | None,
1492
+ limit_name: str,
1493
+ explicit_limit: bool = False,
1494
+ repos: ProxyRepositories,
1495
+ require_fresh_evidence_account_ids: frozenset[str] = frozenset(),
1496
+ ) -> _AdditionalLimitFilterResult:
1497
+ if not accounts:
1498
+ return _AdditionalLimitFilterResult(accounts=[], latest_primary={}, latest_secondary={})
1499
+
1500
+ fresh_since = _additional_usage_fresh_since(to_utc_naive(self._clock.now()))
1501
+ account_ids = [account.id for account in accounts]
1502
+ latest_primary = await _latest_additional_by_key(
1503
+ repos.additional_usage,
1504
+ limit_name,
1505
+ "primary",
1506
+ account_ids=account_ids,
1507
+ )
1508
+ latest_secondary = await _latest_additional_by_key(
1509
+ repos.additional_usage,
1510
+ limit_name,
1511
+ "secondary",
1512
+ account_ids=account_ids,
1513
+ )
1514
+ fresh_primary = await _latest_additional_by_key(
1515
+ repos.additional_usage,
1516
+ limit_name,
1517
+ "primary",
1518
+ account_ids=account_ids,
1519
+ since=fresh_since,
1520
+ )
1521
+ fresh_secondary = await _latest_additional_by_key(
1522
+ repos.additional_usage,
1523
+ limit_name,
1524
+ "secondary",
1525
+ account_ids=account_ids,
1526
+ since=fresh_since,
1527
+ )
1528
+
1529
+ fresh_account_ids = set(fresh_primary) | set(fresh_secondary)
1530
+
1531
+ eligible_accounts: list[Account] = []
1532
+ blocked_by_data = False
1533
+ blocked_by_exhaustion = False
1534
+ now = self._clock.time()
1535
+ for account in accounts:
1536
+ eligibility = _additional_quota_eligibility(
1537
+ account_id=account.id,
1538
+ account_plan_type=account.plan_type,
1539
+ quota_key=limit_name,
1540
+ explicit_limit=explicit_limit,
1541
+ require_fresh_evidence=account.id in require_fresh_evidence_account_ids,
1542
+ latest_primary=latest_primary,
1543
+ latest_secondary=latest_secondary,
1544
+ fresh_primary=fresh_primary,
1545
+ fresh_secondary=fresh_secondary,
1546
+ now=now,
1547
+ )
1548
+ if eligibility == "eligible":
1549
+ eligible_accounts.append(account)
1550
+ continue
1551
+ if eligibility == "data_unavailable":
1552
+ blocked_by_data = True
1553
+ elif eligibility == "quota_exhausted":
1554
+ blocked_by_exhaustion = True
1555
+
1556
+ if not eligible_accounts:
1557
+ if blocked_by_data:
1558
+ error_code = ADDITIONAL_QUOTA_DATA_UNAVAILABLE
1559
+ error_message = f"No fresh additional quota data available for model '{model}'"
1560
+ elif blocked_by_exhaustion:
1561
+ error_code = ADDITIONAL_QUOTA_EXHAUSTED
1562
+ error_message = f"Additional quota exhausted for model '{model}'"
1563
+ else:
1564
+ error_code = NO_ADDITIONAL_QUOTA_ELIGIBLE_ACCOUNTS
1565
+ error_message = f"No accounts with available additional quota for model '{model}'"
1566
+ logger.warning(
1567
+ (
1568
+ "Blocked gated model routing model=%s limit_name=%s reason=%s "
1569
+ "freshness_since=%s candidate_accounts=%s fresh_accounts=%s"
1570
+ ),
1571
+ model,
1572
+ limit_name,
1573
+ error_code,
1574
+ fresh_since.isoformat(),
1575
+ len(accounts),
1576
+ len(fresh_account_ids),
1577
+ )
1578
+ return _AdditionalLimitFilterResult(
1579
+ accounts=[],
1580
+ latest_primary=latest_primary,
1581
+ latest_secondary=latest_secondary,
1582
+ error_code=error_code,
1583
+ error_message=error_message,
1584
+ )
1585
+
1586
+ logger.info(
1587
+ (
1588
+ "Applied gated model routing model=%s limit_name=%s "
1589
+ "candidate_accounts=%s fresh_accounts=%s eligible_accounts=%s"
1590
+ ),
1591
+ model,
1592
+ limit_name,
1593
+ len(accounts),
1594
+ len(fresh_account_ids),
1595
+ len(eligible_accounts),
1596
+ )
1597
+ eligible_ids = {account.id for account in eligible_accounts}
1598
+ return _AdditionalLimitFilterResult(
1599
+ accounts=eligible_accounts,
1600
+ latest_primary={
1601
+ account_id: entry for account_id, entry in latest_primary.items() if account_id in eligible_ids
1602
+ },
1603
+ latest_secondary={
1604
+ account_id: entry for account_id, entry in latest_secondary.items() if account_id in eligible_ids
1605
+ },
1606
+ )
1607
+
1608
+ def _prune_runtime(self, accounts: Iterable[Account]) -> None:
1609
+ account_ids = {account.id for account in accounts}
1610
+ stale_ids = [
1611
+ account_id
1612
+ for account_id, runtime in self._runtime.items()
1613
+ if account_id not in account_ids and not runtime.leases
1614
+ ]
1615
+ for account_id in stale_ids:
1616
+ self._runtime.pop(account_id, None)
1617
+
1618
+ def _prepare_sticky_selection_states(
1619
+ self,
1620
+ selection_inputs: SelectionInputsProtocol,
1621
+ *,
1622
+ required_account_id: str | None,
1623
+ redact_sensitive_details: bool,
1624
+ routing_tunables: RoutingTunables,
1625
+ soft_drain_enabled: bool | None = None,
1626
+ model: str | None = None,
1627
+ ) -> tuple[list[AccountState], dict[str, Account]]:
1628
+ return prepare_selection_states(
1629
+ self,
1630
+ selection_inputs,
1631
+ build_states=_build_states,
1632
+ required_account_id=required_account_id,
1633
+ redact_sensitive_details=redact_sensitive_details,
1634
+ routing_tunables=routing_tunables,
1635
+ soft_drain_enabled=soft_drain_enabled,
1636
+ model=model,
1637
+ )
1638
+
1639
+ async def _get_account_lock(self, account_id: str) -> asyncio.Lock:
1640
+ lock = self._account_locks.get(account_id)
1641
+ if lock is not None:
1642
+ return lock
1643
+ async with self._account_locks_registry_lock:
1644
+ lock = self._account_locks.get(account_id)
1645
+ if lock is None:
1646
+ lock = asyncio.Lock()
1647
+ self._account_locks[account_id] = lock
1648
+ return lock
1649
+
1650
+ async def _select_with_stickiness(
1651
+ self,
1652
+ *,
1653
+ states: list[AccountState],
1654
+ account_map: dict[str, Account],
1655
+ sticky_key: str | None,
1656
+ sticky_kind: StickySessionKind | None,
1657
+ reallocate_sticky: bool,
1658
+ sticky_max_age_seconds: int | None,
1659
+ budget_threshold_pct: float = 95.0,
1660
+ secondary_budget_threshold_pct: float = 100.0,
1661
+ prefer_earlier_reset_accounts: bool,
1662
+ prefer_earlier_reset_window: ResetPreferenceWindow,
1663
+ routing_strategy: RoutingStrategy,
1664
+ relative_availability_power: float = 2.0,
1665
+ relative_availability_top_k: int = 5,
1666
+ sticky_repo: StickySessionsRepository | None,
1667
+ routing_costs_by_account_id: RoutingCostsByAccount | None = None,
1668
+ sticky_existing_account_id: str | None | object = _STICKY_EXISTING_UNSET,
1669
+ initial_preferred_account_id: str | None = None,
1670
+ preserve_existing_mapping_on_fallback: bool = False,
1671
+ preserve_reason_request_local: bool = False,
1672
+ traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND,
1673
+ ignore_standard_quota: bool = False,
1674
+ allow_usage_exhaustion_error: bool = True,
1675
+ usage_exhaustion_states: Iterable[AccountState] | None = None,
1676
+ sticky_refresh_skip_deadline: datetime | None = None,
1677
+ redact_sensitive_details: bool = False,
1678
+ ) -> _StickySelectionOutcome:
1679
+ return await _run_select_with_stickiness(
1680
+ states=states,
1681
+ account_map=account_map,
1682
+ sticky_key=sticky_key,
1683
+ sticky_kind=sticky_kind,
1684
+ reallocate_sticky=reallocate_sticky,
1685
+ sticky_max_age_seconds=sticky_max_age_seconds,
1686
+ budget_threshold_pct=budget_threshold_pct,
1687
+ secondary_budget_threshold_pct=secondary_budget_threshold_pct,
1688
+ prefer_earlier_reset_accounts=prefer_earlier_reset_accounts,
1689
+ prefer_earlier_reset_window=prefer_earlier_reset_window,
1690
+ routing_strategy=routing_strategy,
1691
+ relative_availability_power=relative_availability_power,
1692
+ relative_availability_top_k=relative_availability_top_k,
1693
+ sticky_repo=sticky_repo,
1694
+ routing_costs_by_account_id=routing_costs_by_account_id,
1695
+ sticky_existing_account_id=sticky_existing_account_id,
1696
+ initial_preferred_account_id=initial_preferred_account_id,
1697
+ preserve_existing_mapping_on_fallback=preserve_existing_mapping_on_fallback,
1698
+ preserve_reason_request_local=preserve_reason_request_local,
1699
+ traffic_class=traffic_class,
1700
+ ignore_standard_quota=ignore_standard_quota,
1701
+ allow_usage_exhaustion_error=allow_usage_exhaustion_error,
1702
+ usage_exhaustion_states=usage_exhaustion_states,
1703
+ sticky_refresh_skip_deadline=sticky_refresh_skip_deadline,
1704
+ overload_backoff_runtime=self._runtime,
1705
+ clock=self._clock,
1706
+ redact_sensitive_details=redact_sensitive_details,
1707
+ )
1708
+
1709
+ _persist_sticky_mutation = staticmethod(_persist_sticky_mutation)
1710
+ _restore_sticky_mutation = staticmethod(_restore_sticky_mutation)
1711
+
1712
+ async def mark_rate_limit(self, account: Account, error: UpstreamError) -> None:
1713
+ lock = await self._get_account_lock(account.id)
1714
+ async with lock:
1715
+ state = self._state_for(account)
1716
+ handle_rate_limit(state, error)
1717
+ self._sync_runtime_state(account, state)
1718
+ async with self._repo_factory() as repos:
1719
+ await self._persist_state(repos.accounts, account, state)
1720
+ self._selection_inputs_cache.invalidate()
1721
+
1722
+ async def mark_quota_exceeded(self, account: Account, error: UpstreamError) -> None:
1723
+ lock = await self._get_account_lock(account.id)
1724
+ async with lock:
1725
+ state = self._state_for(account)
1726
+ handle_quota_exceeded(state, error)
1727
+ self._sync_runtime_state(account, state)
1728
+ async with self._repo_factory() as repos:
1729
+ await self._persist_state(repos.accounts, account, state)
1730
+ self._selection_inputs_cache.invalidate()
1731
+
1732
+ async def mark_permanent_failure(self, account: Account, error_code: str) -> bool:
1733
+ """Downgrade *account* to its permanent-failure status.
1734
+
1735
+ Returns whether the permanent downgrade applied (or was already in
1736
+ effect). When the guarded status write MISSES because a peer replica
1737
+ concurrently re-authed/imported and rotated ``refresh_token_encrypted``
1738
+ (the DB row was repaired and left ACTIVE), the account keeps its
1739
+ repaired state. A landed DEACTIVATED downgrade is excluded from local
1740
+ routing; REAUTH_REQUIRED remains request-routable with its stored access
1741
+ token while still blocking future refresh-token exchange.
1742
+ """
1743
+ lock = await self._get_account_lock(account.id)
1744
+ async with lock:
1745
+ state = self._state_for(account)
1746
+ handle_permanent_failure(state, error_code)
1747
+ self._sync_runtime_state(account, state)
1748
+ async with self._repo_factory() as repos:
1749
+ # Guard the DB permanent-status downgrade on the refresh-token
1750
+ # ciphertext this replica currently holds so a concurrent peer
1751
+ # re-auth/import rotation (which changes the ciphertext) is never
1752
+ # clobbered back to a permanent-failure status. On the refresh
1753
+ # path AuthManager._handle_permanent_refresh_failure is the
1754
+ # PRIMARY guarded authority: it has already CAS-written the
1755
+ # downgrade and, in the single-caller case, mutated THIS object's
1756
+ # status to the failure status, so the predicate inside
1757
+ # _persist_state_if_current sees no status change and issues no
1758
+ # redundant write (exactly one guarded downgrade total). This
1759
+ # guarded write covers only the callers whose in-memory object
1760
+ # did not go through that CAS -- an intra-process singleflight
1761
+ # joiner sharing the winner's permanent error, and non-refresh
1762
+ # permanent failures -- without reintroducing the unguarded
1763
+ # update_status that would clobber a peer's ACTIVE/rotated repair
1764
+ # and tear down its live sticky/bridge sessions.
1765
+ downgraded = await self._persist_state_if_current(
1766
+ repos.accounts,
1767
+ account,
1768
+ state,
1769
+ expected_refresh_token_encrypted=account.refresh_token_encrypted,
1770
+ )
1771
+ if downgraded and state.status == AccountStatus.DEACTIVATED:
1772
+ mark_account_routing_unavailable(account.id)
1773
+ self._selection_inputs_cache.invalidate()
1774
+ return downgraded
1775
+
1776
+ async def record_error(self, account: Account) -> None:
1777
+ await self.record_errors(account, 1)
1778
+
1779
+ async def record_error_backoff(self, account: Account) -> None:
1780
+ """Record one error and immediately enter the bounded transient backoff."""
1781
+ await self.record_errors(account, 1, minimum_error_count=ERROR_BACKOFF_THRESHOLD)
1782
+
1783
+ async def record_errors(
1784
+ self,
1785
+ account: Account,
1786
+ count: int,
1787
+ *,
1788
+ minimum_error_count: int = 0,
1789
+ ) -> None:
1790
+ """Record *count* transient errors in a single lock acquisition."""
1791
+ if count < 1:
1792
+ return
1793
+ lock = await self._get_account_lock(account.id)
1794
+ async with lock:
1795
+ account_snapshot = _clone_account(account)
1796
+ state = self._state_for(account)
1797
+ state.error_count = max(state.error_count + count, minimum_error_count)
1798
+ state.last_error_at = self._clock.time()
1799
+ self._sync_runtime_state(account, state)
1800
+ runtime = self._runtime.get(account.id)
1801
+ record_outcome_locked(self._runtime[account.id], state.last_error_at, success=False, count=count)
1802
+ if runtime and runtime.health_tier == HEALTH_TIER_PROBING:
1803
+ runtime.probe_success_streak = 0
1804
+ async with self._repo_factory() as repos:
1805
+ await self._persist_state_if_current(repos.accounts, account_snapshot, state)
1806
+
1807
+ async def record_success(self, account: Account) -> None:
1808
+ """Clear transient error state after a successful upstream request."""
1809
+ lock = await self._get_account_lock(account.id)
1810
+ async with lock:
1811
+ runtime = self._runtime.setdefault(account.id, RuntimeState())
1812
+ record_outcome_locked(runtime, self._clock.time(), success=True)
1813
+ if runtime and runtime.error_count > 0:
1814
+ runtime.error_count = 0
1815
+ runtime.last_error_at = None
1816
+ runtime.version += 1
1817
+ runtime.health_version += 1
1818
+ if runtime and runtime.health_tier == HEALTH_TIER_PROBING:
1819
+ runtime.probe_success_streak += 1
1820
+ runtime.version += 1
1821
+ runtime.health_version += 1
1822
+
1823
+ async def record_probe_result(
1824
+ self,
1825
+ *,
1826
+ account_id: str,
1827
+ http_status: int,
1828
+ ) -> None:
1829
+ """Settle an operator probe into this replica's advisory health state."""
1830
+ lock = await self._get_account_lock(account_id)
1831
+ if not 200 <= http_status < 300:
1832
+ async with lock:
1833
+ runtime = self._runtime.setdefault(account_id, RuntimeState())
1834
+ if runtime.probe_success_streak > 0:
1835
+ runtime.probe_success_streak = 0
1836
+ runtime.version += 1
1837
+ runtime.health_version += 1
1838
+ return
1839
+
1840
+ # Usage reads intentionally run without the per-account lock. Capture a
1841
+ # health-observation token first so an older successful probe cannot
1842
+ # clear a failure or health-tier change recorded while those reads are
1843
+ # in flight, while lease-only pressure changes remain harmless.
1844
+ async with lock:
1845
+ expected_health_version = self._runtime.setdefault(account_id, RuntimeState()).health_version
1846
+
1847
+ async with self._repo_factory() as repos:
1848
+ account = await repos.accounts.get_by_id(account_id)
1849
+ if account is None:
1850
+ return
1851
+ primary_entry = await repos.usage.latest_entry_for_account(account_id, window="primary")
1852
+ secondary_entry = await repos.usage.latest_entry_for_account(account_id, window="secondary")
1853
+ monthly_entry = await repos.usage.latest_entry_for_account(account_id, window="monthly")
1854
+ # Force Probe must interpret refreshed rows exactly like ordinary
1855
+ # routing: raw storage slots do not identify weekly/monthly meaning.
1856
+ effective_secondary_entry = _select_long_window_entry(
1857
+ account=account,
1858
+ monthly_entry=monthly_entry,
1859
+ secondary_entry=secondary_entry,
1860
+ )
1861
+ # Repository teardown rolls back reads and expires ORM rows. Keep
1862
+ # loaded snapshots for normalization after the session has closed.
1863
+ account = _clone_account(account)
1864
+ primary_entry = clone_row(primary_entry) if primary_entry is not None else None
1865
+ effective_secondary_entry = (
1866
+ clone_row(effective_secondary_entry) if effective_secondary_entry is not None else None
1867
+ )
1868
+ now = self._clock.time()
1869
+ normalized_usage = _normalize_usage_inputs(
1870
+ account=account,
1871
+ primary_entry=primary_entry,
1872
+ secondary_entry=effective_secondary_entry,
1873
+ now_epoch=int(now),
1874
+ )
1875
+ health_primary_used = _health_tier_primary_used(
1876
+ plan_type=account.plan_type,
1877
+ primary_used=normalized_usage.primary_used,
1878
+ )
1879
+ routing_policy = _normalize_account_routing_policy(account.routing_policy)
1880
+ # C2-3 resilience toggles: one dashboard snapshot before the lock.
1881
+ resilience = resolve_resilience_toggles(await get_settings_cache().get())
1882
+
1883
+ # C2-2 routing/overload: resolved before the lock, the knobs selection uses.
1884
+ tunables = self.current_routing_tunables()
1885
+ async with lock:
1886
+ runtime = self._runtime.setdefault(account_id, RuntimeState())
1887
+ # Treat settlement as a local CAS: the newer runtime health
1888
+ # observation wins, and a later probe can retry with a fresh usage
1889
+ # snapshot. Lease-only version bumps must not drop Force Probe.
1890
+ if runtime.health_version != expected_health_version:
1891
+ return
1892
+
1893
+ normalized_state = _state_from_account(
1894
+ account=account,
1895
+ primary_entry=primary_entry,
1896
+ secondary_entry=effective_secondary_entry,
1897
+ runtime=replace(runtime),
1898
+ routing_tunables=tunables,
1899
+ now=now,
1900
+ soft_drain_enabled=resilience.soft_drain_enabled,
1901
+ )
1902
+ account_status = normalized_state.status
1903
+ if account_status not in (AccountStatus.ACTIVE, AccountStatus.REAUTH_REQUIRED):
1904
+ return
1905
+
1906
+ was_probe_eligible = runtime.health_tier == HEALTH_TIER_PROBING
1907
+ if was_probe_eligible and (runtime.error_count > 0 or runtime.last_error_at is not None):
1908
+ runtime.error_count = 0
1909
+ runtime.last_error_at = None
1910
+ runtime.version += 1
1911
+ runtime.health_version += 1
1912
+
1913
+ _sync_runtime_health_tier(
1914
+ account_id=account_id,
1915
+ status=account_status,
1916
+ used_percent=health_primary_used,
1917
+ secondary_used_percent=normalized_usage.secondary_used,
1918
+ routing_policy=routing_policy,
1919
+ runtime=runtime,
1920
+ now=now,
1921
+ soft_drain_enabled=resilience.soft_drain_enabled,
1922
+ )
1923
+ if runtime.health_tier != HEALTH_TIER_PROBING:
1924
+ return
1925
+ if not was_probe_eligible and (runtime.error_count > 0 or runtime.last_error_at is not None):
1926
+ runtime.error_count = 0
1927
+ runtime.last_error_at = None
1928
+ runtime.version += 1
1929
+ runtime.health_version += 1
1930
+
1931
+ runtime.probe_success_streak += 1
1932
+ runtime.version += 1
1933
+ runtime.health_version += 1
1934
+ _sync_runtime_health_tier(
1935
+ account_id=account_id,
1936
+ status=account_status,
1937
+ used_percent=health_primary_used,
1938
+ secondary_used_percent=normalized_usage.secondary_used,
1939
+ routing_policy=routing_policy,
1940
+ runtime=runtime,
1941
+ now=now,
1942
+ soft_drain_enabled=resilience.soft_drain_enabled,
1943
+ )
1944
+
1945
+ def _state_for(self, account: Account) -> AccountState:
1946
+ runtime = self._runtime.setdefault(account.id, RuntimeState())
1947
+ routing_policy = _normalize_account_routing_policy(getattr(account, "routing_policy", None))
1948
+ return AccountState(
1949
+ account_id=account.id,
1950
+ status=account.status,
1951
+ used_percent=None,
1952
+ reset_at=runtime.reset_at,
1953
+ primary_reset_at=None,
1954
+ blocked_at=float(account.blocked_at) if account.blocked_at is not None else runtime.blocked_at,
1955
+ cooldown_until=runtime.cooldown_until,
1956
+ secondary_used_percent=None,
1957
+ secondary_reset_at=None,
1958
+ last_error_at=runtime.last_error_at,
1959
+ last_selected_at=runtime.last_selected_at,
1960
+ error_count=runtime.error_count,
1961
+ deactivation_reason=account.deactivation_reason,
1962
+ plan_type=account.plan_type,
1963
+ capacity_credits=usage_core.capacity_for_plan(account.plan_type, "secondary"),
1964
+ routing_policy=routing_policy,
1965
+ ignore_standard_quota=False,
1966
+ )
1967
+
1968
+ def _sync_runtime_state(
1969
+ self,
1970
+ account: Account,
1971
+ state: AccountState,
1972
+ *,
1973
+ selected: bool = False,
1974
+ expected_version: int | None = None,
1975
+ ) -> bool:
1976
+ runtime = self._runtime.setdefault(account.id, RuntimeState())
1977
+ if expected_version is not None and runtime.version != expected_version:
1978
+ if selected:
1979
+ runtime.last_selected_at = self._clock.time()
1980
+ runtime.version += 1
1981
+ return False
1982
+
1983
+ dirty = False
1984
+ if runtime.reset_at != state.reset_at:
1985
+ runtime.reset_at = state.reset_at
1986
+ dirty = True
1987
+ if runtime.cooldown_until != state.cooldown_until:
1988
+ runtime.cooldown_until = state.cooldown_until
1989
+ dirty = True
1990
+ if runtime.blocked_at != state.blocked_at:
1991
+ runtime.blocked_at = state.blocked_at
1992
+ dirty = True
1993
+ if runtime.last_error_at != state.last_error_at:
1994
+ runtime.last_error_at = state.last_error_at
1995
+ dirty = True
1996
+ if runtime.error_count != state.error_count:
1997
+ runtime.error_count = state.error_count
1998
+ dirty = True
1999
+ if account.status != state.status:
2000
+ dirty = True
2001
+ if account.deactivation_reason != state.deactivation_reason:
2002
+ dirty = True
2003
+ health_dirty = dirty
2004
+ if selected:
2005
+ runtime.last_selected_at = self._clock.time()
2006
+ dirty = True
2007
+ if dirty:
2008
+ runtime.version += 1
2009
+ if health_dirty:
2010
+ runtime.health_version += 1
2011
+ return True
2012
+
2013
+ async def _persist_selection_state(
2014
+ self,
2015
+ accounts_repo: AccountsRepository,
2016
+ account_map: dict[str, Account],
2017
+ states: list[AccountState],
2018
+ ) -> set[str]:
2019
+ stale_account_ids: set[str] = set()
2020
+ for state in states:
2021
+ if state.ignore_standard_quota:
2022
+ continue
2023
+ account = account_map.get(state.account_id)
2024
+ if account is not None:
2025
+ persisted = await self._persist_state_if_current(accounts_repo, account, state)
2026
+ if not persisted:
2027
+ stale_account_ids.add(account.id)
2028
+ return stale_account_ids
2029
+
2030
+ async def _persist_state(
2031
+ self,
2032
+ accounts_repo: AccountsRepository,
2033
+ account: Account,
2034
+ state: AccountState,
2035
+ ) -> None:
2036
+ reset_at_int = int(state.reset_at) if state.reset_at else None
2037
+ blocked_at_int = int(state.blocked_at) if state.blocked_at else None
2038
+ status_changed = account.status != state.status
2039
+ reason_changed = account.deactivation_reason != state.deactivation_reason
2040
+ reset_changed = account.reset_at != reset_at_int
2041
+ blocked_changed = account.blocked_at != blocked_at_int
2042
+
2043
+ if status_changed or reason_changed or reset_changed or blocked_changed:
2044
+ await accounts_repo.update_status(
2045
+ account.id,
2046
+ state.status,
2047
+ state.deactivation_reason,
2048
+ reset_at_int,
2049
+ blocked_at=blocked_at_int,
2050
+ )
2051
+ account.status = state.status
2052
+ account.deactivation_reason = state.deactivation_reason
2053
+ account.reset_at = reset_at_int
2054
+ account.blocked_at = blocked_at_int
2055
+
2056
+ async def _persist_state_if_current(
2057
+ self,
2058
+ accounts_repo: AccountsRepository,
2059
+ account: Account,
2060
+ state: AccountState,
2061
+ *,
2062
+ expected_refresh_token_encrypted: bytes | None = None,
2063
+ ) -> bool:
2064
+ reset_at_int = int(state.reset_at) if state.reset_at else None
2065
+ blocked_at_int = int(state.blocked_at) if state.blocked_at else None
2066
+ status_changed = account.status != state.status
2067
+ reason_changed = account.deactivation_reason != state.deactivation_reason
2068
+ reset_changed = account.reset_at != reset_at_int
2069
+ blocked_changed = account.blocked_at != blocked_at_int
2070
+
2071
+ if status_changed or reason_changed or reset_changed or blocked_changed:
2072
+ updated = await accounts_repo.update_status_if_current(
2073
+ account.id,
2074
+ state.status,
2075
+ state.deactivation_reason,
2076
+ reset_at_int,
2077
+ blocked_at=blocked_at_int,
2078
+ expected_status=account.status,
2079
+ expected_deactivation_reason=account.deactivation_reason,
2080
+ expected_reset_at=account.reset_at,
2081
+ expected_blocked_at=account.blocked_at,
2082
+ expected_refresh_token_encrypted=expected_refresh_token_encrypted,
2083
+ )
2084
+ if updated:
2085
+ account.status = state.status
2086
+ account.deactivation_reason = state.deactivation_reason
2087
+ account.reset_at = reset_at_int
2088
+ account.blocked_at = blocked_at_int
2089
+ return updated
2090
+ return True
2091
+
2092
+ async def _sync_state(
2093
+ self,
2094
+ accounts_repo: AccountsRepository,
2095
+ account: Account,
2096
+ state: AccountState,
2097
+ ) -> None:
2098
+ self._sync_runtime_state(account, state)
2099
+ await self._persist_state(accounts_repo, account, state)
2100
+
2101
+
2102
+ def _build_states(
2103
+ *,
2104
+ accounts: Iterable[Account],
2105
+ latest_primary: Mapping[str, UsageHistory | AdditionalUsageHistory],
2106
+ latest_secondary: Mapping[str, UsageHistory | AdditionalUsageHistory],
2107
+ latest_monthly: Mapping[str, UsageHistory],
2108
+ runtime: dict[str, RuntimeState],
2109
+ now: float | None = None,
2110
+ routing_policy_override: str | None = None,
2111
+ ignore_standard_quota_account_ids: frozenset[str] = frozenset(),
2112
+ encryptor: TokenEncryptor | None = None,
2113
+ routing_tunables: RoutingTunables | None = None,
2114
+ soft_drain_enabled: bool | None = None,
2115
+ model: str | None = None,
2116
+ log_weight_transitions: bool = True,
2117
+ ) -> tuple[list[AccountState], dict[str, Account]]:
2118
+ now = REAL_CLOCK.time() if now is None else now
2119
+ # Request and background callers pass their snapshot's values; None (tests, tools) = environment layer.
2120
+ tunables = routing_tunables or effective_routing_tunables()
2121
+ states: list[AccountState] = []
2122
+ account_map: dict[str, Account] = {}
2123
+
2124
+ for account in accounts:
2125
+ secondary_entry: UsageHistory | AdditionalUsageHistory | None = latest_secondary.get(account.id)
2126
+ if account.id not in ignore_standard_quota_account_ids:
2127
+ secondary_entry = _select_long_window_entry(
2128
+ account=account,
2129
+ monthly_entry=latest_monthly.get(account.id),
2130
+ secondary_entry=secondary_entry,
2131
+ )
2132
+ state = _state_from_account(
2133
+ account=account,
2134
+ primary_entry=latest_primary.get(account.id),
2135
+ secondary_entry=secondary_entry,
2136
+ runtime=runtime.setdefault(account.id, RuntimeState()),
2137
+ access_token_expires_at=(
2138
+ account_access_token_expires_at(account, encryptor)
2139
+ if account.status == AccountStatus.REAUTH_REQUIRED and encryptor is not None
2140
+ else None
2141
+ ),
2142
+ soft_drain_enabled=soft_drain_enabled,
2143
+ now=now,
2144
+ routing_tunables=tunables,
2145
+ )
2146
+ if routing_policy_override is not None and account.id in ignore_standard_quota_account_ids:
2147
+ state.routing_policy = routing_policy_override
2148
+ state.ignore_standard_quota = account.id in ignore_standard_quota_account_ids
2149
+ states.append(state)
2150
+ account_map[account.id] = account
2151
+ apply_latency_cohort_weights(states, runtime, now=now, model=model, log_transitions=log_weight_transitions)
2152
+ return states, account_map
2153
+
2154
+
2155
+ def effective_routing_tunables(dashboard_settings: object | None = None) -> RoutingTunables:
2156
+ """Routing/overload knobs from a cached dashboard row; ``None`` = environment alone (C2-2)."""
2157
+ return resolve_routing_tunables(dashboard_settings, startup_settings=get_settings())
2158
+
2159
+
2160
+ def effective_account_concurrency_caps(dashboard_settings: object | None = None) -> AccountConcurrencyCaps:
2161
+ startup_settings = get_settings()
2162
+ configured_response_create_limit, configured_stream_limit = configured_account_concurrency_caps(
2163
+ dashboard_settings, startup_settings=startup_settings
2164
+ )
2165
+ scope = getattr(startup_settings, "proxy_account_caps_scope", "partitioned")
2166
+ partition = get_cap_partition()
2167
+ if scope == "replica" or partition.replica_count <= 1:
2168
+ return AccountConcurrencyCaps(
2169
+ response_create_limit=configured_response_create_limit,
2170
+ stream_limit=configured_stream_limit,
2171
+ )
2172
+ return AccountConcurrencyCaps(
2173
+ response_create_limit=partition_cap(configured_response_create_limit, partition.replica_count, partition.rank),
2174
+ stream_limit=partition_cap(configured_stream_limit, partition.replica_count, partition.rank),
2175
+ configured_response_create_limit=configured_response_create_limit,
2176
+ configured_stream_limit=configured_stream_limit,
2177
+ replica_count=partition.replica_count,
2178
+ )
2179
+
2180
+
2181
+ def _record_account_lease_acquired(kind: AccountLeaseKind) -> None:
2182
+ if PROMETHEUS_AVAILABLE and account_lease_acquired_total is not None:
2183
+ account_lease_acquired_total.labels(kind=kind).inc()
2184
+
2185
+
2186
+ def _record_account_lease_released(kind: AccountLeaseKind, reason: str) -> None:
2187
+ if PROMETHEUS_AVAILABLE and account_lease_released_total is not None:
2188
+ account_lease_released_total.labels(kind=kind, reason=reason).inc()
2189
+
2190
+
2191
+ def _record_account_lease_stale_reclaimed(kind: AccountLeaseKind) -> None:
2192
+ if PROMETHEUS_AVAILABLE and account_lease_stale_reclaimed_total is not None:
2193
+ account_lease_stale_reclaimed_total.labels(kind=kind).inc()
2194
+
2195
+
2196
+ def _record_account_inflight_leases(account_id: str, runtime: RuntimeState) -> None:
2197
+ if PROMETHEUS_AVAILABLE and account_inflight_leases is not None:
2198
+ account_inflight_leases.labels(account_id=account_id, kind="response_create").set(
2199
+ runtime.inflight_response_creates
2200
+ )
2201
+ account_inflight_leases.labels(account_id=account_id, kind="stream").set(runtime.inflight_streams)
2202
+
2203
+
2204
+ def _record_account_cap_rejection(kind: AccountLeaseKind | None) -> None:
2205
+ if kind is None:
2206
+ return
2207
+ if PROMETHEUS_AVAILABLE and account_cap_rejections_total is not None:
2208
+ account_cap_rejections_total.labels(kind=kind).inc()
2209
+
2210
+
2211
+ def _record_stream_pool_gauges(pool_capacity: int, pool_inflight: int) -> None:
2212
+ if PROMETHEUS_AVAILABLE and stream_pool_capacity is not None:
2213
+ stream_pool_capacity.set(pool_capacity)
2214
+ if PROMETHEUS_AVAILABLE and stream_pool_inflight is not None:
2215
+ stream_pool_inflight.set(pool_inflight)
2216
+
2217
+
2218
+ def _record_api_key_fair_share_rejection() -> None:
2219
+ if PROMETHEUS_AVAILABLE and api_key_fair_share_rejections_total is not None:
2220
+ api_key_fair_share_rejections_total.inc()
2221
+
2222
+
2223
+ def _normalize_account_routing_policy(value: str | None) -> str:
2224
+ if value in _ACCOUNT_ROUTING_POLICIES:
2225
+ return value
2226
+ return _ROUTING_POLICY_NORMAL
2227
+
2228
+
2229
+ async def _load_dashboard_additional_quota_routing_overrides() -> dict[str, str]:
2230
+ dashboard_settings = await get_settings_cache().get()
2231
+ return _parse_additional_quota_routing_policies(dashboard_settings.additional_quota_routing_policies_json)
2232
+
2233
+
2234
+ def _additional_quota_routing_policy_override(limit_name: str | None, policies: dict[str, str]) -> str | None:
2235
+ if limit_name is None:
2236
+ return None
2237
+ normalized_limit_name = canonicalize_additional_quota_key(limit_name=limit_name)
2238
+ if normalized_limit_name is None:
2239
+ return None
2240
+ policy = get_additional_quota_routing_policy(normalized_limit_name, overrides=policies)
2241
+ if policy == "inherit":
2242
+ return None
2243
+ return policy
2244
+
2245
+
2246
+ def _parse_additional_quota_routing_policies(raw_policies: str) -> dict[str, str]:
2247
+ if not raw_policies:
2248
+ return {}
2249
+ try:
2250
+ parsed = json.loads(raw_policies)
2251
+ except json.JSONDecodeError:
2252
+ return {}
2253
+ if not isinstance(parsed, dict):
2254
+ return {}
2255
+ policies: dict[str, str] = {}
2256
+ for quota_key, policy in parsed.items():
2257
+ if not isinstance(quota_key, str) or not isinstance(policy, str):
2258
+ continue
2259
+ normalized_key = normalize_additional_quota_key(quota_key)
2260
+ normalized_policy = policy.strip().lower()
2261
+ if normalized_key and normalized_policy in _ADDITIONAL_QUOTA_ROUTING_POLICIES:
2262
+ policies[normalized_key] = normalized_policy
2263
+ return policies
2264
+
2265
+
2266
+ def _state_from_account(
2267
+ *,
2268
+ account: Account,
2269
+ primary_entry: UsageHistory | AdditionalUsageHistory | None,
2270
+ secondary_entry: UsageHistory | AdditionalUsageHistory | None,
2271
+ runtime: RuntimeState,
2272
+ access_token_expires_at: float | None = None,
2273
+ now: float | None = None,
2274
+ routing_tunables: RoutingTunables | None = None,
2275
+ soft_drain_enabled: bool | None = None,
2276
+ ) -> AccountState:
2277
+ now = REAL_CLOCK.time() if now is None else now
2278
+ tunables = routing_tunables or effective_routing_tunables()
2279
+ routing_policy = _normalize_account_routing_policy(getattr(account, "routing_policy", None))
2280
+ normalized_usage = _normalize_usage_inputs(
2281
+ account=account,
2282
+ primary_entry=primary_entry,
2283
+ secondary_entry=secondary_entry,
2284
+ now_epoch=int(now),
2285
+ )
2286
+ primary_used = normalized_usage.primary_used
2287
+ primary_reset = normalized_usage.primary_reset
2288
+ primary_window_minutes = normalized_usage.primary_window_minutes
2289
+ effective_secondary_entry = normalized_usage.effective_secondary_entry
2290
+ secondary_used = normalized_usage.secondary_used
2291
+ secondary_reset = normalized_usage.secondary_reset
2292
+ effective_blocked_at = float(account.blocked_at) if account.blocked_at is not None else runtime.blocked_at
2293
+ credits_has, credits_unlimited, credits_balance = _extract_credit_status(
2294
+ primary_entry,
2295
+ effective_secondary_entry,
2296
+ secondary_entry,
2297
+ recorded_after=effective_blocked_at if account.status == AccountStatus.QUOTA_EXCEEDED else None,
2298
+ )
2299
+
2300
+ # If the usage window has reset (reset_at is in the past), the last
2301
+ # recorded sample describes an expired window at ANY used percentage:
2302
+ # upstream may have stopped reporting the window entirely (e.g. the
2303
+ # temporary 5h-limit removal), in which case the row is never rewritten
2304
+ # and a frozen sub-100% sample would otherwise hold drain tiers and
2305
+ # budget pressure forever. Zero the derived locals — not the stored
2306
+ # rows — so the account is not incorrectly blocked or deprioritised
2307
+ # while waiting for the next usage refresh. Expired samples map to 0.0
2308
+ # rather than None because usage-derived status recovery only evaluates
2309
+ # non-None percentages.
2310
+ if primary_used is not None and primary_reset is not None and primary_reset <= int(now):
2311
+ primary_used = 0.0
2312
+ primary_reset = None
2313
+ # A strictly newer long-window row proves a later fetch no longer
2314
+ # reported the short window: drop the stale duration — whether or not
2315
+ # the stale row's reset has elapsed — so phase planning stops treating
2316
+ # the account as having a short phase window.
2317
+ if (
2318
+ primary_window_minutes is not None
2319
+ and primary_entry is not None
2320
+ and effective_secondary_entry is not None
2321
+ and effective_secondary_entry is not primary_entry
2322
+ and (effective_secondary_entry.recorded_at - primary_entry.recorded_at).total_seconds()
2323
+ > _SIBLING_FETCH_MARGIN_SECONDS
2324
+ ):
2325
+ primary_window_minutes = None
2326
+ if secondary_used is not None and secondary_reset is not None and secondary_reset <= int(now):
2327
+ secondary_used = 0.0
2328
+ secondary_reset = None
2329
+ ignore_zero_capacity_primary_runtime_reset = False
2330
+ status_seed = account.status
2331
+ long_window_quota_available = _usage_entry_is_recent_available(effective_secondary_entry, now=now)
2332
+ # Preserve actual 429 cooldowns (marked by blocked_at) before zero-primary
2333
+ # recovery can rewrite status from fresh long-window usage.
2334
+ rate_limited_cooldown_deadline: float | None = None
2335
+ if account.status == AccountStatus.RATE_LIMITED and effective_blocked_at is not None:
2336
+ persisted_deadline = plausible_rate_limit_reset_at(account.reset_at, now=now) or (
2337
+ effective_blocked_at + RATE_LIMITED_MIN_COOLDOWN_SECONDS
2338
+ )
2339
+ if now < persisted_deadline:
2340
+ rate_limited_cooldown_deadline = persisted_deadline
2341
+ if (
2342
+ rate_limited_cooldown_deadline is not None
2343
+ and runtime.cooldown_until is not None
2344
+ and runtime.cooldown_until <= now
2345
+ and runtime.blocked_at is not None
2346
+ and runtime.blocked_at >= effective_blocked_at
2347
+ ):
2348
+ # Only the replica that observed this block may recover before its persisted deadline.
2349
+ early_freshness_entry = _rate_limited_freshness_entry(
2350
+ account=account,
2351
+ primary_entry=primary_entry,
2352
+ long_window_entry=effective_secondary_entry,
2353
+ now=now,
2354
+ )
2355
+ if _usage_entry_recorded_after_block(early_freshness_entry, effective_blocked_at):
2356
+ rate_limited_cooldown_deadline = None
2357
+
2358
+ if usage_core.capacity_for_plan(account.plan_type, "primary") == 0.0 and (
2359
+ account.status != AccountStatus.RATE_LIMITED
2360
+ or (
2361
+ rate_limited_cooldown_deadline is None
2362
+ and (
2363
+ (
2364
+ primary_window_minutes is not None
2365
+ and not usage_core.is_primary_window_minutes(primary_window_minutes)
2366
+ and long_window_quota_available
2367
+ )
2368
+ or (primary_entry is None and long_window_quota_available)
2369
+ )
2370
+ )
2371
+ ):
2372
+ primary_used = _health_tier_primary_used(
2373
+ plan_type=account.plan_type,
2374
+ primary_used=primary_used,
2375
+ )
2376
+ primary_reset = None
2377
+ primary_window_minutes = None
2378
+ ignore_zero_capacity_primary_runtime_reset = account.status == AccountStatus.RATE_LIMITED
2379
+ if account.status == AccountStatus.RATE_LIMITED:
2380
+ status_seed = AccountStatus.ACTIVE
2381
+
2382
+ # Use account.reset_at from DB as the authoritative source for runtime reset
2383
+ # and to survive process restarts.
2384
+ persisted_reset_at = float(account.reset_at) if account.reset_at is not None else None
2385
+ runtime_reset_at = runtime.reset_at
2386
+ # Validate only future RATE_LIMITED hints. Elapsed deadlines must still
2387
+ # reach apply_usage_quota's ordinary expiry transition, and QUOTA_EXCEEDED
2388
+ # deadlines have separate recovery semantics.
2389
+ if account.status == AccountStatus.RATE_LIMITED:
2390
+ if persisted_reset_at is not None and persisted_reset_at > now:
2391
+ persisted_reset_at = plausible_rate_limit_reset_at(persisted_reset_at, now=now)
2392
+ if runtime_reset_at is not None and runtime_reset_at > now:
2393
+ runtime_reset_at = plausible_rate_limit_reset_at(runtime_reset_at, now=now)
2394
+ rejected_persisted_rate_limit_reset = (
2395
+ account.status == AccountStatus.RATE_LIMITED
2396
+ and account.reset_at is not None
2397
+ and persisted_reset_at is None
2398
+ and account.reset_at > now
2399
+ )
2400
+ db_reset_at = None if ignore_zero_capacity_primary_runtime_reset else persisted_reset_at
2401
+ if status_seed in (AccountStatus.RATE_LIMITED, AccountStatus.QUOTA_EXCEEDED) or runtime.blocked_at is not None:
2402
+ effective_runtime_reset = db_reset_at or runtime_reset_at
2403
+ else:
2404
+ effective_runtime_reset = None
2405
+
2406
+ # Resetless rate limits retain a minimum hold after blocked_at across restarts;
2407
+ # after this floor, recovery uses the normal compare-and-set persistence path.
2408
+ if (
2409
+ status_seed == AccountStatus.RATE_LIMITED
2410
+ and effective_runtime_reset is None
2411
+ and effective_blocked_at is not None
2412
+ ):
2413
+ floor_deadline = effective_blocked_at + RATE_LIMITED_MIN_COOLDOWN_SECONDS
2414
+ if now < floor_deadline:
2415
+ effective_runtime_reset = floor_deadline
2416
+
2417
+ if (
2418
+ account.status == AccountStatus.QUOTA_EXCEEDED
2419
+ and effective_runtime_reset is not None
2420
+ and effective_runtime_reset > now
2421
+ and effective_blocked_at is None
2422
+ and effective_secondary_entry is not None
2423
+ and long_window_quota_available
2424
+ and effective_secondary_entry.reset_at is not None
2425
+ and float(effective_secondary_entry.reset_at) > effective_runtime_reset
2426
+ ):
2427
+ effective_runtime_reset = None
2428
+
2429
+ # Post-block evidence clears resets after debounce. Quota recovery uses persisted
2430
+ # markers; early rate-limit recovery requires this replica's runtime block evidence.
2431
+ cooldown_ready = False
2432
+ if account.status == AccountStatus.QUOTA_EXCEEDED:
2433
+ cooldown_ready = (
2434
+ effective_blocked_at is not None and now >= effective_blocked_at + QUOTA_EXCEEDED_COOLDOWN_SECONDS
2435
+ )
2436
+ elif (
2437
+ runtime.cooldown_until is not None
2438
+ and runtime.cooldown_until <= now
2439
+ and runtime.blocked_at is not None
2440
+ and effective_blocked_at is not None
2441
+ and runtime.blocked_at >= effective_blocked_at
2442
+ ):
2443
+ cooldown_ready = True
2444
+
2445
+ if cooldown_ready and effective_blocked_at is not None:
2446
+ if account.status == AccountStatus.QUOTA_EXCEEDED:
2447
+ freshness_entry = (
2448
+ effective_secondary_entry if secondary_used is not None and secondary_used < 100.0 else None
2449
+ )
2450
+ elif account.status == AccountStatus.RATE_LIMITED:
2451
+ freshness_entry = _rate_limited_freshness_entry(
2452
+ account=account,
2453
+ primary_entry=primary_entry,
2454
+ long_window_entry=effective_secondary_entry,
2455
+ now=now,
2456
+ )
2457
+ else:
2458
+ freshness_entry = None
2459
+ if _usage_entry_recorded_after_block(freshness_entry, effective_blocked_at):
2460
+ effective_runtime_reset = None
2461
+
2462
+ rejected_reset_recovery_evidence = False
2463
+ if rejected_persisted_rate_limit_reset:
2464
+ rejected_reset_freshness_entry = _rate_limited_freshness_entry(
2465
+ account=account,
2466
+ primary_entry=primary_entry,
2467
+ long_window_entry=effective_secondary_entry,
2468
+ now=now,
2469
+ )
2470
+ # Recovery requires actual evidence without exhaustion in any applicable window.
2471
+ all_quota_windows_available = (
2472
+ (primary_used is None or float(primary_used) < 100.0)
2473
+ and (secondary_used is None or float(secondary_used) < 100.0)
2474
+ and (primary_used is not None or secondary_used is not None)
2475
+ )
2476
+ rejected_reset_recovery_evidence = all_quota_windows_available and _usage_entry_is_recent_available(
2477
+ rejected_reset_freshness_entry, now=now
2478
+ )
2479
+ if effective_blocked_at is not None:
2480
+ # A sample predating the 429 cannot disprove the persisted block.
2481
+ rejected_reset_recovery_evidence = (
2482
+ rejected_reset_recovery_evidence
2483
+ and now >= effective_blocked_at + RATE_LIMITED_MIN_COOLDOWN_SECONDS
2484
+ and _usage_entry_recorded_after_block(rejected_reset_freshness_entry, effective_blocked_at)
2485
+ )
2486
+
2487
+ # A resetless rate limit whose runtime cooldown was lost (e.g. a restart
2488
+ # after a 429 without reset metadata) has no deadline to expire and no
2489
+ # post-block evidence trail; a long-window sample alone must not clear
2490
+ # it. Evidence-gated clearing above always starts from a persisted or
2491
+ # runtime reset, so this only matches the truly resetless case.
2492
+ resetless_rate_limit_without_evidence = (
2493
+ status_seed == AccountStatus.RATE_LIMITED and account.reset_at is None and runtime.reset_at is None
2494
+ )
2495
+
2496
+ quota_secondary_used = secondary_used
2497
+ if (
2498
+ status_seed == AccountStatus.QUOTA_EXCEEDED
2499
+ and secondary_used is not None
2500
+ and secondary_used >= 100.0
2501
+ and effective_secondary_entry is not None
2502
+ and (
2503
+ not _usage_entry_is_recent_enough(effective_secondary_entry.recorded_at, now=now)
2504
+ or (
2505
+ effective_blocked_at is not None
2506
+ and not _usage_entry_recorded_after_block(effective_secondary_entry, effective_blocked_at)
2507
+ )
2508
+ )
2509
+ ):
2510
+ # Historical exhaustion cannot rewrite a newer upstream rejection's deadline.
2511
+ quota_secondary_used = None
2512
+
2513
+ status, used_percent, reset_at = apply_usage_quota(
2514
+ status=status_seed,
2515
+ primary_used=primary_used,
2516
+ primary_reset=primary_reset,
2517
+ primary_window_minutes=primary_window_minutes,
2518
+ runtime_reset=effective_runtime_reset,
2519
+ secondary_used=quota_secondary_used,
2520
+ secondary_reset=secondary_reset,
2521
+ credits_has=credits_has,
2522
+ credits_unlimited=credits_unlimited,
2523
+ credits_balance=credits_balance,
2524
+ infer_status_from_usage=False,
2525
+ now=now,
2526
+ )
2527
+ if resetless_rate_limit_without_evidence and primary_used is None and status == AccountStatus.ACTIVE:
2528
+ status = AccountStatus.RATE_LIMITED
2529
+ if rejected_persisted_rate_limit_reset and not rejected_reset_recovery_evidence:
2530
+ status = AccountStatus.RATE_LIMITED
2531
+ reset_at = float(account.reset_at)
2532
+
2533
+ next_blocked_at = (
2534
+ effective_blocked_at
2535
+ if status == AccountStatus.QUOTA_EXCEEDED
2536
+ or (status == AccountStatus.RATE_LIMITED and account.status != AccountStatus.QUOTA_EXCEEDED)
2537
+ else None
2538
+ )
2539
+
2540
+ if soft_drain_enabled is None:
2541
+ # C2-3 resilience toggles: callers pass the dashboard value; None (tests, tools) = env alias / default.
2542
+ soft_drain_enabled = resolve_resilience_toggles(None, startup_settings=get_settings()).soft_drain_enabled
2543
+ new_tier = _sync_runtime_health_tier(
2544
+ account_id=account.id,
2545
+ status=status,
2546
+ used_percent=used_percent,
2547
+ secondary_used_percent=secondary_used,
2548
+ routing_policy=routing_policy,
2549
+ runtime=runtime,
2550
+ now=now,
2551
+ soft_drain_enabled=soft_drain_enabled,
2552
+ )
2553
+
2554
+ inflight_pressure_pct = (
2555
+ runtime.inflight_response_creates + runtime.inflight_streams
2556
+ ) * tunables.inflight_penalty_pct
2557
+ leased_token_pressure_pct = 0.0
2558
+ long_window_key = "secondary"
2559
+ if effective_secondary_entry is not None and effective_secondary_entry.window == "monthly":
2560
+ long_window_key = "monthly"
2561
+ capacity_credits = usage_core.capacity_for_plan(account.plan_type, long_window_key) or 0.0
2562
+ if capacity_credits > 0.0 and runtime.leased_tokens > 0:
2563
+ leased_token_pressure_pct = runtime.leased_tokens * tunables.lease_token_weight / capacity_credits * 100.0
2564
+ pressure_pct = inflight_pressure_pct + leased_token_pressure_pct
2565
+ effective_used_percent = None if used_percent is None else min(100.0, used_percent + pressure_pct)
2566
+ effective_secondary_used_percent = None if secondary_used is None else min(100.0, secondary_used + pressure_pct)
2567
+ usage_exhaustion_evidence_status = status in (AccountStatus.QUOTA_EXCEEDED, AccountStatus.RATE_LIMITED)
2568
+
2569
+ return AccountState(
2570
+ account_id=account.id,
2571
+ status=status,
2572
+ used_percent=effective_used_percent,
2573
+ reset_at=reset_at,
2574
+ primary_reset_at=primary_reset,
2575
+ primary_window_minutes=primary_window_minutes,
2576
+ blocked_at=next_blocked_at,
2577
+ cooldown_until=runtime.cooldown_until,
2578
+ secondary_used_percent=effective_secondary_used_percent,
2579
+ secondary_reset_at=secondary_reset,
2580
+ last_error_at=runtime.last_error_at,
2581
+ last_selected_at=runtime.last_selected_at,
2582
+ error_count=runtime.error_count,
2583
+ deactivation_reason=account.deactivation_reason,
2584
+ plan_type=account.plan_type,
2585
+ capacity_credits=capacity_credits,
2586
+ health_tier=new_tier,
2587
+ priority_used_percent=used_percent if usage_exhaustion_evidence_status else None,
2588
+ priority_secondary_used_percent=secondary_used if usage_exhaustion_evidence_status else None,
2589
+ access_token_expires_at=access_token_expires_at,
2590
+ inflight_response_creates=runtime.inflight_response_creates,
2591
+ inflight_streams=runtime.inflight_streams,
2592
+ leased_tokens=runtime.leased_tokens,
2593
+ routing_policy=routing_policy,
2594
+ selection_weight_multiplier=error_rate_weight_multiplier(
2595
+ runtime, now, policy=ErrorRateWeightingPolicy(enabled=tunables.error_rate_weighting_enabled)
2596
+ ),
2597
+ )
2598
+
2599
+
2600
+ def _normalize_usage_inputs(
2601
+ *,
2602
+ account: Account,
2603
+ primary_entry: _UsageWindowEntry | None,
2604
+ secondary_entry: _UsageWindowEntry | None,
2605
+ now_epoch: int,
2606
+ ) -> _NormalizedUsageInputs:
2607
+ """Normalize persisted usage for routing and explicit probe settlement."""
2608
+ primary_used = primary_entry.used_percent if primary_entry else None
2609
+ primary_reset = primary_entry.reset_at if primary_entry else None
2610
+ primary_window_minutes = primary_entry.window_minutes if primary_entry else None
2611
+ effective_secondary_entry = secondary_entry
2612
+ if (
2613
+ effective_secondary_entry is not None
2614
+ and effective_secondary_entry.window == "monthly"
2615
+ and usage_core.capacity_for_plan(account.plan_type, "monthly") is None
2616
+ ):
2617
+ effective_secondary_entry = None
2618
+ primary_row = usage_history_to_window_row(primary_entry) if primary_entry is not None else None
2619
+ secondary_row = usage_history_to_window_row(secondary_entry) if secondary_entry is not None else None
2620
+ # Weekly-only accounts may not emit a dedicated secondary row; treat the
2621
+ # weekly primary row as quota-window input for balancer decisions. When
2622
+ # both rows exist, prefer the newer weekly snapshot.
2623
+ if primary_row is not None and usage_core.should_use_weekly_primary(primary_row, secondary_row):
2624
+ effective_secondary_entry = primary_entry
2625
+ primary_used = None
2626
+ primary_reset = None
2627
+ primary_window_minutes = None
2628
+
2629
+ secondary_used = effective_secondary_entry.used_percent if effective_secondary_entry else None
2630
+ secondary_reset = effective_secondary_entry.reset_at if effective_secondary_entry else None
2631
+
2632
+ # Expired rows describe prior windows. Zero derived values without
2633
+ # rewriting history so stale samples cannot hold drain tiers forever.
2634
+ if primary_used is not None and primary_reset is not None and primary_reset <= now_epoch:
2635
+ primary_used = 0.0
2636
+ primary_reset = None
2637
+ # A strictly newer long-window row proves a later fetch no longer
2638
+ # reported the short window, so phase planning drops the stale duration.
2639
+ if (
2640
+ primary_window_minutes is not None
2641
+ and primary_entry is not None
2642
+ and effective_secondary_entry is not None
2643
+ and effective_secondary_entry is not primary_entry
2644
+ and (effective_secondary_entry.recorded_at - primary_entry.recorded_at).total_seconds()
2645
+ > _SIBLING_FETCH_MARGIN_SECONDS
2646
+ ):
2647
+ primary_window_minutes = None
2648
+ if secondary_used is not None and secondary_reset is not None and secondary_reset <= now_epoch:
2649
+ secondary_used = 0.0
2650
+ secondary_reset = None
2651
+
2652
+ return _NormalizedUsageInputs(
2653
+ primary_used=primary_used,
2654
+ primary_reset=primary_reset,
2655
+ primary_window_minutes=primary_window_minutes,
2656
+ effective_secondary_entry=effective_secondary_entry,
2657
+ secondary_used=secondary_used,
2658
+ secondary_reset=secondary_reset,
2659
+ )
2660
+
2661
+
2662
+ def _health_tier_primary_used(*, plan_type: str | None, primary_used: float | None) -> float | None:
2663
+ """Drop primary usage when the plan has no primary-window capacity."""
2664
+ # Storage may retain a legacy/synthetic primary row for free accounts. The
2665
+ # health state machine must follow plan capacity, not the row's slot, or
2666
+ # both ordinary routing and Force Probe can drain an account on a quota it
2667
+ # does not have.
2668
+ if usage_core.capacity_for_plan(plan_type, "primary") == 0.0:
2669
+ return None
2670
+ return primary_used
2671
+
2672
+
2673
+ def _sync_runtime_health_tier(
2674
+ *,
2675
+ account_id: str,
2676
+ status: AccountStatus,
2677
+ used_percent: float | None,
2678
+ secondary_used_percent: float | None,
2679
+ routing_policy: str,
2680
+ runtime: RuntimeState,
2681
+ now: float,
2682
+ soft_drain_enabled: bool,
2683
+ ) -> int:
2684
+ before = (
2685
+ runtime.health_tier,
2686
+ runtime.drain_entered_at,
2687
+ runtime.probe_success_streak,
2688
+ )
2689
+ if soft_drain_enabled:
2690
+ new_tier = evaluate_health_tier(
2691
+ AccountState(
2692
+ account_id=account_id,
2693
+ status=status,
2694
+ used_percent=used_percent,
2695
+ secondary_used_percent=secondary_used_percent,
2696
+ last_error_at=runtime.last_error_at,
2697
+ error_count=runtime.error_count,
2698
+ health_tier=runtime.health_tier,
2699
+ routing_policy=routing_policy,
2700
+ ),
2701
+ now=now,
2702
+ drain_entered_at=runtime.drain_entered_at,
2703
+ probe_success_streak=runtime.probe_success_streak,
2704
+ # Drain/probe thresholds are fixed in
2705
+ # ``app/core/balancer/logic.py`` (evaluate_health_tier defaults).
2706
+ )
2707
+ if new_tier == HEALTH_TIER_DRAINING and runtime.health_tier != HEALTH_TIER_DRAINING:
2708
+ runtime.drain_entered_at = now
2709
+ runtime.probe_success_streak = 0
2710
+ if new_tier == HEALTH_TIER_HEALTHY:
2711
+ runtime.drain_entered_at = None
2712
+ runtime.probe_success_streak = 0
2713
+ runtime.health_tier = new_tier
2714
+ else:
2715
+ runtime.health_tier = HEALTH_TIER_HEALTHY
2716
+ runtime.drain_entered_at = None
2717
+ runtime.probe_success_streak = 0
2718
+
2719
+ after = (
2720
+ runtime.health_tier,
2721
+ runtime.drain_entered_at,
2722
+ runtime.probe_success_streak,
2723
+ )
2724
+ if after != before:
2725
+ runtime.version += 1
2726
+ runtime.health_version += 1
2727
+ return runtime.health_tier
2728
+
2729
+
2730
+ def background_recovery_state_from_account(
2731
+ *,
2732
+ account: Account,
2733
+ primary_entry: UsageHistory | None,
2734
+ secondary_entry: UsageHistory | None,
2735
+ routing_tunables: RoutingTunables | None = None,
2736
+ soft_drain_enabled: bool | None = None,
2737
+ ) -> AccountState:
2738
+ """Evaluate recovery without live runtime state.
2739
+
2740
+ Seed a throwaway runtime from the persisted block marker so post-block usage
2741
+ can clear stale reset guards after a balancer restart. ``routing_tunables`` and
2742
+ ``soft_drain_enabled`` are the caller's dashboard-snapshot values (health tier follows the dashboard).
2743
+ """
2744
+
2745
+ runtime = RuntimeState()
2746
+ blocked_at = float(account.blocked_at) if account.blocked_at is not None else None
2747
+ now = REAL_CLOCK.time()
2748
+ reset_at = float(account.reset_at) if account.reset_at is not None else None
2749
+ valid_reset_at = plausible_rate_limit_reset_at(reset_at, now=now)
2750
+
2751
+ if blocked_at is not None:
2752
+ runtime.blocked_at = blocked_at
2753
+
2754
+ if account.status == AccountStatus.RATE_LIMITED and blocked_at is not None and valid_reset_at is not None:
2755
+ runtime.cooldown_until = valid_reset_at
2756
+ state = _state_from_account(
2757
+ account=account,
2758
+ primary_entry=primary_entry,
2759
+ secondary_entry=secondary_entry,
2760
+ runtime=runtime,
2761
+ now=now,
2762
+ routing_tunables=routing_tunables,
2763
+ soft_drain_enabled=soft_drain_enabled,
2764
+ )
2765
+ if account.status == AccountStatus.RATE_LIMITED:
2766
+ freshness_entry = _rate_limited_freshness_entry(
2767
+ account=account,
2768
+ primary_entry=primary_entry,
2769
+ long_window_entry=secondary_entry,
2770
+ now=now,
2771
+ )
2772
+ # Keep elapsed resets intact until _state_from_account evaluates the
2773
+ # selector's normal expiry path; only freshness gates the final repair.
2774
+ if blocked_at is not None and reset_at is not None and reset_at <= now:
2775
+ minimum_floor_deadline = blocked_at + RATE_LIMITED_MIN_COOLDOWN_SECONDS
2776
+ # An early explicit reset does not let scheduler reconciliation
2777
+ # bypass the persisted post-429 minimum floor.
2778
+ if now < minimum_floor_deadline or not _usage_entry_recorded_after_block(freshness_entry, blocked_at):
2779
+ return replace(
2780
+ state,
2781
+ status=AccountStatus.RATE_LIMITED,
2782
+ reset_at=reset_at,
2783
+ blocked_at=blocked_at,
2784
+ cooldown_until=max(reset_at, minimum_floor_deadline),
2785
+ )
2786
+ elif blocked_at is None and reset_at is not None and reset_at <= now:
2787
+ if not _usage_entry_is_recent_available(freshness_entry, now=now):
2788
+ return replace(
2789
+ state,
2790
+ status=AccountStatus.RATE_LIMITED,
2791
+ reset_at=reset_at,
2792
+ blocked_at=None,
2793
+ cooldown_until=None,
2794
+ )
2795
+ if reset_at is None:
2796
+ return replace(
2797
+ state,
2798
+ status=AccountStatus.RATE_LIMITED,
2799
+ reset_at=None,
2800
+ blocked_at=blocked_at,
2801
+ cooldown_until=None,
2802
+ )
2803
+ return state
2804
+
2805
+
2806
+ def _select_long_window_entry(
2807
+ *,
2808
+ account: Account,
2809
+ monthly_entry: UsageHistory | None,
2810
+ secondary_entry: UsageHistory | AdditionalUsageHistory | None,
2811
+ ) -> UsageHistory | AdditionalUsageHistory | None:
2812
+ if monthly_entry is not None and usage_core.capacity_for_plan(account.plan_type, "monthly") is not None:
2813
+ return monthly_entry
2814
+ return secondary_entry
2815
+
2816
+
2817
+ def _rate_limited_freshness_entry(
2818
+ *,
2819
+ account: Account,
2820
+ primary_entry: _UsageWindowEntry | None,
2821
+ long_window_entry: _UsageWindowEntry | None,
2822
+ now: float,
2823
+ ) -> _UsageWindowEntry | None:
2824
+ if (
2825
+ long_window_entry is not None
2826
+ and long_window_entry.reset_at is not None
2827
+ and long_window_entry.reset_at <= int(now)
2828
+ ):
2829
+ long_window_entry = None
2830
+ if (
2831
+ long_window_entry is not None
2832
+ and long_window_entry.window == "monthly"
2833
+ and usage_core.capacity_for_plan(account.plan_type, "monthly") is None
2834
+ ):
2835
+ long_window_entry = None
2836
+ # Freshness cannot prove recovery while an applicable long window is
2837
+ # still exhausted, even if the primary sample reports available quota.
2838
+ if long_window_entry is not None and not (
2839
+ long_window_entry.used_percent is not None and float(long_window_entry.used_percent) < 100.0
2840
+ ):
2841
+ return None
2842
+ if long_window_entry is not None and long_window_entry.window == "monthly":
2843
+ return long_window_entry
2844
+ if primary_entry is None:
2845
+ return long_window_entry
2846
+ # A newer long-window row can replace primary evidence only after the
2847
+ # primary reset expires. Otherwise the primary sample must itself
2848
+ # report available quota.
2849
+ primary_window_expired = primary_entry.reset_at is not None and float(primary_entry.reset_at) <= now
2850
+ if (
2851
+ primary_window_expired
2852
+ and long_window_entry is not None
2853
+ and long_window_entry.recorded_at > primary_entry.recorded_at
2854
+ ):
2855
+ return long_window_entry
2856
+ if primary_entry.used_percent is not None and float(primary_entry.used_percent) < 100.0:
2857
+ return primary_entry
2858
+ return None
2859
+
2860
+
2861
+ def _usage_entry_is_recent_available(entry: _UsageWindowEntry | None, *, now: float) -> bool:
2862
+ return (
2863
+ entry is not None
2864
+ and _usage_entry_is_recent_enough(entry.recorded_at, now=now)
2865
+ and entry.used_percent is not None
2866
+ and float(entry.used_percent) < 100.0
2867
+ )
2868
+
2869
+
2870
+ def _usage_entry_recorded_after_block(entry: _UsageWindowEntry | None, blocked_at: float) -> bool:
2871
+ if entry is None or entry.recorded_at is None:
2872
+ return False
2873
+ recorded_at = entry.recorded_at
2874
+ if recorded_at.tzinfo is None:
2875
+ recorded_at = recorded_at.replace(tzinfo=timezone.utc)
2876
+ # Persistence truncates block timestamps to whole seconds. A sample
2877
+ # within that same second cannot prove it was captured after the block.
2878
+ return int(recorded_at.timestamp()) > int(blocked_at)
2879
+
2880
+
2881
+ def _extract_credit_status(
2882
+ *entries: _UsageWindowEntry | None,
2883
+ recorded_after: float | None = None,
2884
+ ) -> tuple[bool | None, bool | None, float | None]:
2885
+ credit_entries: list[UsageHistory] = [
2886
+ entry
2887
+ for entry in entries
2888
+ if isinstance(entry, UsageHistory)
2889
+ and (recorded_after is None or _usage_entry_recorded_after_block(entry, recorded_after))
2890
+ and not (entry.credits_has is None and entry.credits_unlimited is None and entry.credits_balance is None)
2891
+ ]
2892
+ if not credit_entries:
2893
+ return None, None, None
2894
+ entry = max(
2895
+ credit_entries,
2896
+ key=lambda item: item.recorded_at if item.recorded_at is not None else datetime.min,
2897
+ )
2898
+ if entry is not None:
2899
+ return entry.credits_has, entry.credits_unlimited, entry.credits_balance
2900
+ return None, None, None
2901
+
2902
+
2903
+ def _usage_entry_is_recent_enough(recorded_at: datetime | None, *, now: float) -> bool:
2904
+ if recorded_at is None:
2905
+ return False
2906
+ current_time = datetime.fromtimestamp(now, tz=timezone.utc)
2907
+ interval_seconds = usage_freshness_horizon_seconds()
2908
+ recorded_time = recorded_at if recorded_at.tzinfo is not None else recorded_at.replace(tzinfo=timezone.utc)
2909
+ return recorded_time >= current_time - timedelta(seconds=interval_seconds)
2910
+
2911
+
2912
+ def _filter_accounts_for_model(
2913
+ accounts: list[Account],
2914
+ model: str,
2915
+ *,
2916
+ service_tier: str | None = None,
2917
+ ) -> list[Account]:
2918
+ return _filter_accounts_for_model_impl(
2919
+ accounts,
2920
+ model,
2921
+ registry=get_model_registry(),
2922
+ service_tier=service_tier,
2923
+ )
2924
+
2925
+
2926
+ def _filter_accounts_for_model_with_catalog_evidence(
2927
+ accounts: list[Account],
2928
+ model: str,
2929
+ *,
2930
+ service_tier: str | None = None,
2931
+ additional_quota_can_override_account_catalog: bool = False,
2932
+ ) -> _ModelAccountFilterResult:
2933
+ return _filter_accounts_for_model_with_catalog_evidence_impl(
2934
+ accounts,
2935
+ model,
2936
+ registry=get_model_registry(),
2937
+ service_tier=service_tier,
2938
+ additional_quota_can_override_account_catalog=additional_quota_can_override_account_catalog,
2939
+ )
2940
+
2941
+
2942
+ def _selectable_accounts(accounts: list[Account]) -> list[Account]:
2943
+ return [account for account in accounts if account.status not in (AccountStatus.DEACTIVATED, AccountStatus.PAUSED)]
2944
+
2945
+
2946
+ def _mapped_model_has_registry_entry(model: str | None) -> bool:
2947
+ return _mapped_model_has_registry_entry_impl(model, registry=get_model_registry())
2948
+
2949
+
2950
+ def _first_not_none(
2951
+ primary_entry: UsageHistory | AdditionalUsageHistory | None,
2952
+ secondary_entry: UsageHistory | AdditionalUsageHistory | None,
2953
+ field: str,
2954
+ ):
2955
+ if primary_entry is not None:
2956
+ value = getattr(primary_entry, field, None)
2957
+ if value is not None:
2958
+ return value
2959
+ if secondary_entry is not None:
2960
+ return getattr(secondary_entry, field, None)
2961
+ return None
2962
+
2963
+
2964
+ def _clone_selection_inputs(selection_inputs: SelectionInputs) -> SelectionInputs:
2965
+ return _SelectionInputs(
2966
+ accounts=[_clone_account(account) for account in selection_inputs.accounts],
2967
+ latest_primary={account_id: clone_row(entry) for account_id, entry in selection_inputs.latest_primary.items()},
2968
+ latest_secondary={
2969
+ account_id: clone_row(entry) for account_id, entry in selection_inputs.latest_secondary.items()
2970
+ },
2971
+ latest_monthly={account_id: clone_row(entry) for account_id, entry in selection_inputs.latest_monthly.items()},
2972
+ continuity_owner_candidates=(
2973
+ None
2974
+ if selection_inputs.continuity_owner_candidates is None
2975
+ else [_clone_account(account) for account in selection_inputs.continuity_owner_candidates]
2976
+ ),
2977
+ sticky_mutation_authority_account_ids=(
2978
+ None
2979
+ if selection_inputs.sticky_mutation_authority_account_ids is None
2980
+ else frozenset(selection_inputs.sticky_mutation_authority_account_ids)
2981
+ ),
2982
+ quota_planner_settings=selection_inputs.quota_planner_settings,
2983
+ runtime_accounts=(
2984
+ None
2985
+ if selection_inputs.runtime_accounts is None
2986
+ else [_clone_account(account) for account in selection_inputs.runtime_accounts]
2987
+ ),
2988
+ error_message=selection_inputs.error_message,
2989
+ error_code=selection_inputs.error_code,
2990
+ ignore_standard_quota_account_ids=frozenset(selection_inputs.ignore_standard_quota_account_ids),
2991
+ ignore_standard_quota_status=selection_inputs.ignore_standard_quota_status,
2992
+ persist_standard_quota_status=selection_inputs.persist_standard_quota_status,
2993
+ routing_policy_override=selection_inputs.routing_policy_override,
2994
+ quota_admitted_catalog_omission_account_ids=frozenset(
2995
+ selection_inputs.quota_admitted_catalog_omission_account_ids
2996
+ ),
2997
+ )
2998
+
2999
+
3000
+ def _additional_usage_fresh_since(now: datetime | None = None) -> datetime:
3001
+ current_time = now or utcnow()
3002
+ return current_time - timedelta(seconds=usage_freshness_horizon_seconds())
3003
+
3004
+
3005
+ def _is_upstream_circuit_breaker_open(circuit_breaker_enabled: bool) -> bool:
3006
+ # C2-3 resilience toggles: the caller resolved the flag from its snapshot.
3007
+ if not circuit_breaker_enabled:
3008
+ return False
3009
+ return are_all_account_circuit_breakers_open()
3010
+
3011
+
3012
+ def _format_degraded_error_message(message: str | None) -> str:
3013
+ degradation_status = get_degradation_status()
3014
+ reason = degradation_status.get("reason") or "upstream capacity is currently unavailable"
3015
+ base_message = message or "Upstream unavailable"
3016
+ return f"{base_message}. Service is operating in degraded mode: {reason}"