Skip to main content

ironflow_store/entities/
provider_account.rs

1//! Provider Account entities: AI provider accounts, their usage windows and history.
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use strum::{Display, EnumString};
6use uuid::Uuid;
7
8/// Prefix of the system secrets holding Provider Account credentials.
9///
10/// Keys under this prefix are hidden from the Secrets listing and refused by
11/// the Secrets API.
12pub const PROVIDER_ACCOUNT_SECRET_PREFIX: &str = "accounts/";
13
14/// Secret key holding the credential of the account `id`.
15///
16/// # Examples
17///
18/// ```
19/// use ironflow_store::entities::provider_account_secret_key;
20/// use uuid::Uuid;
21///
22/// let key = provider_account_secret_key(Uuid::nil());
23/// assert_eq!(key, "accounts/00000000-0000-0000-0000-000000000000/credential");
24/// ```
25pub fn provider_account_secret_key(id: Uuid) -> String {
26    format!("{PROVIDER_ACCOUNT_SECRET_PREFIX}{id}/credential")
27}
28
29/// Status of a usage window, as reported by the provider.
30///
31/// # Examples
32///
33/// ```
34/// use ironflow_store::entities::AccountWindowStatus;
35///
36/// let status: AccountWindowStatus = "allowed_warning".parse().unwrap();
37/// assert_eq!(status, AccountWindowStatus::AllowedWarning);
38/// ```
39#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
40#[cfg_attr(feature = "store-postgres", derive(sqlx::Type))]
41#[cfg_attr(
42    feature = "store-postgres",
43    sqlx(type_name = "text", rename_all = "snake_case")
44)]
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Display, EnumString)]
46#[serde(rename_all = "snake_case")]
47#[strum(serialize_all = "snake_case")]
48pub enum AccountWindowStatus {
49    /// Requests are allowed.
50    Allowed,
51    /// Requests are allowed, close to the limit.
52    AllowedWarning,
53    /// Requests are rejected until the window resets.
54    Rejected,
55}
56
57/// A persisted Provider Account.
58///
59/// The credential itself lives in the secret `secret_key`, never here.
60#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
61pub struct ProviderAccount {
62    /// Account ID (UUID v7).
63    pub id: Uuid,
64    /// Unique slug.
65    pub name: String,
66    /// Human-readable name.
67    pub display_name: String,
68    /// Account kind (e.g. `claude_subscription`).
69    pub kind: String,
70    /// Key of the secret holding the credential.
71    pub secret_key: String,
72    /// Whether the account may be selected.
73    pub enabled: bool,
74    /// Lower values are preferred by the `priority` strategy.
75    pub priority: i32,
76    /// Free-form tags.
77    pub tags: Vec<String>,
78    /// Maximum concurrent steps, `None` for unlimited.
79    pub max_concurrency: Option<u32>,
80    /// Utilization from which the account is shown as near its limit.
81    pub alert_threshold: f64,
82    /// When the credential expires.
83    pub expires_at: DateTime<Utc>,
84    /// Subscription plan (`pro`, `max`), informative.
85    pub plan: Option<String>,
86    /// When the provider last rejected the credential, `None` when valid.
87    pub auth_failed_at: Option<DateTime<Utc>>,
88    /// User who created the account.
89    pub created_by: Option<Uuid>,
90    /// Creation timestamp.
91    pub created_at: DateTime<Utc>,
92    /// Last update timestamp.
93    pub updated_at: DateTime<Utc>,
94}
95
96/// Request to create a Provider Account.
97///
98/// The caller generates `id` so the secret key can be built first.
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct NewProviderAccount {
101    /// Account ID (UUID v7).
102    pub id: Uuid,
103    /// Unique slug.
104    pub name: String,
105    /// Human-readable name.
106    pub display_name: String,
107    /// Account kind.
108    pub kind: String,
109    /// Key of the secret holding the credential.
110    pub secret_key: String,
111    /// Whether the account may be selected.
112    pub enabled: bool,
113    /// Priority.
114    pub priority: i32,
115    /// Tags.
116    pub tags: Vec<String>,
117    /// Maximum concurrent steps.
118    pub max_concurrency: Option<u32>,
119    /// Alert threshold in `(0, 1]`.
120    pub alert_threshold: f64,
121    /// When the credential expires.
122    pub expires_at: DateTime<Utc>,
123    /// Subscription plan.
124    pub plan: Option<String>,
125    /// Creating user.
126    pub created_by: Option<Uuid>,
127}
128
129/// Partial update of a Provider Account. `None` fields are left unchanged.
130///
131/// # Examples
132///
133/// ```
134/// use ironflow_store::entities::ProviderAccountUpdate;
135///
136/// let update = ProviderAccountUpdate {
137///     enabled: Some(false),
138///     max_concurrency: Some(None),
139///     ..ProviderAccountUpdate::default()
140/// };
141/// assert_eq!(update.max_concurrency, Some(None));
142/// ```
143#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
144pub struct ProviderAccountUpdate {
145    /// New display name.
146    pub display_name: Option<String>,
147    /// Enable or disable.
148    pub enabled: Option<bool>,
149    /// New priority.
150    pub priority: Option<i32>,
151    /// New tags (replaces).
152    pub tags: Option<Vec<String>>,
153    /// New max concurrency; `Some(None)` clears it.
154    pub max_concurrency: Option<Option<u32>>,
155    /// New alert threshold.
156    pub alert_threshold: Option<f64>,
157    /// New expiry.
158    pub expires_at: Option<DateTime<Utc>>,
159    /// New plan; `Some(None)` clears it.
160    pub plan: Option<Option<String>>,
161    /// New auth failure mark; `Some(None)` clears it.
162    pub auth_failed_at: Option<Option<DateTime<Utc>>>,
163}
164
165/// Latest reading of one usage window of an account.
166#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
167#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
168pub struct ProviderAccountWindow {
169    /// Account the window belongs to.
170    pub account_id: Uuid,
171    /// Window name (`five_hour`, `seven_day`).
172    pub window: String,
173    /// Fraction used, `0.0..=1.0`.
174    pub utilization: f64,
175    /// When the window resets.
176    pub resets_at: Option<DateTime<Utc>>,
177    /// Provider status.
178    pub status: AccountWindowStatus,
179    /// Model family the window applies to, `None` for every model.
180    pub model_scope: Option<String>,
181    /// When the window was observed.
182    pub observed_at: DateTime<Utc>,
183}
184
185/// One historical observation of a window.
186#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
187pub struct ProviderAccountUsagePoint {
188    /// Observation ID.
189    pub id: Uuid,
190    /// Account the observation belongs to.
191    pub account_id: Uuid,
192    /// Window name.
193    pub window: String,
194    /// Fraction used.
195    pub utilization: f64,
196    /// When the window resets.
197    pub resets_at: Option<DateTime<Utc>>,
198    /// Provider status.
199    pub status: AccountWindowStatus,
200    /// Model scope.
201    pub model_scope: Option<String>,
202    /// When the window was observed.
203    pub observed_at: DateTime<Utc>,
204}
205
206/// A window observed during an invocation or a credential check.
207#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
208pub struct NewAccountWindow {
209    /// Window name.
210    pub window: String,
211    /// Fraction used.
212    pub utilization: f64,
213    /// When the window resets.
214    pub resets_at: Option<DateTime<Utc>>,
215    /// Provider status.
216    pub status: AccountWindowStatus,
217    /// Model scope.
218    pub model_scope: Option<String>,
219    /// When the window was observed.
220    pub observed_at: DateTime<Utc>,
221}
222
223/// Observations to record for an account.
224///
225/// # Examples
226///
227/// ```
228/// use ironflow_store::entities::NewProviderAccountObservation;
229///
230/// let observation = NewProviderAccountObservation { windows: Vec::new(), auth_failed: true };
231/// assert!(observation.auth_failed);
232/// ```
233#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
234pub struct NewProviderAccountObservation {
235    /// Observed windows.
236    pub windows: Vec<NewAccountWindow>,
237    /// Whether the provider rejected the credential.
238    #[serde(default)]
239    pub auth_failed: bool,
240}
241
242/// An account eligible for selection, with its windows and load.
243#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
244pub struct ProviderAccountCandidate {
245    /// The account.
246    pub account: ProviderAccount,
247    /// Its latest windows.
248    pub windows: Vec<ProviderAccountWindow>,
249    /// Steps currently running under it.
250    pub running_steps: u32,
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    #[test]
258    fn secret_key_is_under_the_accounts_prefix() {
259        let id = Uuid::now_v7();
260        let key = provider_account_secret_key(id);
261        assert!(key.starts_with(PROVIDER_ACCOUNT_SECRET_PREFIX));
262        assert!(key.ends_with("/credential"));
263        assert!(key.contains(&id.to_string()));
264    }
265
266    #[test]
267    fn window_status_wire_format() {
268        assert_eq!(AccountWindowStatus::Rejected.to_string(), "rejected");
269        assert_eq!(
270            serde_json::to_string(&AccountWindowStatus::AllowedWarning).unwrap(),
271            "\"allowed_warning\""
272        );
273        assert!("nope".parse::<AccountWindowStatus>().is_err());
274    }
275
276    #[test]
277    fn observation_auth_failed_defaults_to_false() {
278        let observation: NewProviderAccountObservation =
279            serde_json::from_str(r#"{"windows":[]}"#).unwrap();
280        assert!(!observation.auth_failed);
281    }
282}