Skip to main content

artifactize_core/
limits.rs

1//! Machine-wide limits in `$STATE/limits.json`: `backends` caps how many Agent reviews of a
2//! backend are in flight at once across every `verify` that shares the state directory, and
3//! `agentSessions` sets whether Agent conversations are saved and how large their store grows.
4//!
5//! ```json
6//! {"backends": {"codex": 4, "openai": 16},
7//!  "agentSessions": {"enabled": true, "maxBytes": 1073741824, "targetBytes": 805306368}}
8//! ```
9
10use std::{collections::BTreeMap, path::Path};
11
12use serde::{Deserialize, Serialize};
13
14use crate::config::Backend;
15#[cfg(test)]
16mod backend_tests;
17
18pub const FILE: &str = "limits.json";
19/// Machine-wide limits are a small declaration; reject oversized files before JSON allocation.
20const MAX_BYTES: u64 = 64 * 1024;
21/// Reject implausible capacity declarations while allowing large hosts; absent backends
22/// remain unlimited rather than being silently clamped to this declared-slot bound.
23const MAX_SLOTS: u32 = 100_000;
24/// The session store's default size that starts a collection: 1 GiB.
25pub const SESSIONS_MAX_BYTES: u64 = 1024 * 1024 * 1024;
26/// The size a collection brings the session store down to by default: 768 MiB.
27pub const SESSIONS_TARGET_BYTES: u64 = 768 * 1024 * 1024;
28/// Collect to three quarters of a small declared maximum to leave growth headroom.
29/// Divide first, preserving quarter-block rounding and avoiding multiplication overflow.
30const SESSION_TARGET_NUMERATOR: u64 = 3;
31const SESSION_TARGET_DENOMINATOR: u64 = 4;
32
33/// Backend names to their machine-wide slot counts; a backend without an entry is unlimited.
34#[derive(Debug, Clone, Default, PartialEq, Eq)]
35pub struct Limits {
36    backends: BTreeMap<Backend, u32>,
37    agent_sessions: AgentSessions,
38}
39
40/// Saving Agent conversations and the size bounds of their store. When the store grows past
41/// `max_bytes`, the oldest sessions are deleted until it is at or below `target_bytes`.
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
43#[serde(rename_all = "camelCase")]
44pub struct AgentSessions {
45    pub enabled: bool,
46    pub max_bytes: u64,
47    pub target_bytes: u64,
48}
49
50impl Default for AgentSessions {
51    fn default() -> Self {
52        Self {
53            enabled: true,
54            max_bytes: SESSIONS_MAX_BYTES,
55            target_bytes: SESSIONS_TARGET_BYTES,
56        }
57    }
58}
59
60#[derive(Deserialize)]
61#[serde(deny_unknown_fields)]
62struct Declared {
63    #[serde(default, deserialize_with = "backend_limits")]
64    backends: BTreeMap<Backend, u32>,
65    #[serde(default, rename = "agentSessions")]
66    agent_sessions: Option<DeclaredSessions>,
67}
68
69fn backend_limits<'de, D: serde::Deserializer<'de>>(
70    deserializer: D,
71) -> Result<BTreeMap<Backend, u32>, D::Error> {
72    // Parse each external name once, retain its Backend, and preserve existing diagnostics
73    // and lexicographic validation order at this declaration boundary.
74    let names = BTreeMap::<String, u32>::deserialize(deserializer)?;
75    names
76        .into_iter()
77        .map(|(name, limit)| {
78            let backend =
79                serde_json::from_value::<Backend>(serde_json::Value::String(name.clone()))
80                    .map_err(|error| serde::de::Error::custom(format!("backends: {error}")))?;
81            if !(1..=MAX_SLOTS).contains(&limit) {
82                return Err(serde::de::Error::custom(format!(
83                    "backends.{name} must be between 1 and {MAX_SLOTS}."
84                )));
85            }
86            Ok((backend, limit))
87        })
88        .collect()
89}
90
91#[derive(Deserialize)]
92#[serde(rename_all = "camelCase", deny_unknown_fields)]
93struct DeclaredSessions {
94    enabled: Option<bool>,
95    max_bytes: Option<u64>,
96    target_bytes: Option<u64>,
97}
98
99impl DeclaredSessions {
100    /// The declared bounds over the defaults. Without `targetBytes`, a `maxBytes` at or below
101    /// the 768 MiB default collects down to three quarters of it.
102    fn resolve(self) -> Result<AgentSessions, String> {
103        let max_bytes = self.max_bytes.unwrap_or(SESSIONS_MAX_BYTES);
104        if max_bytes == 0 {
105            return Err("agentSessions.maxBytes must be at least 1.".into());
106        }
107        let target_bytes = match self.target_bytes {
108            Some(target) if target >= max_bytes => {
109                return Err(format!(
110                    "agentSessions.targetBytes ({target}) must be lower than maxBytes ({max_bytes})."
111                ));
112            }
113            Some(target) => target,
114            None if SESSIONS_TARGET_BYTES < max_bytes => SESSIONS_TARGET_BYTES,
115            None => max_bytes / SESSION_TARGET_DENOMINATOR * SESSION_TARGET_NUMERATOR,
116        };
117        Ok(AgentSessions {
118            enabled: self.enabled.unwrap_or(true),
119            max_bytes,
120            target_bytes,
121        })
122    }
123}
124
125impl Limits {
126    /// Read `$STATE/limits.json`. No file means no limits; a file that is not a regular file,
127    /// is not valid, names an unknown backend or a limit outside 1–100000, or declares
128    /// `agentSessions` bounds out of order is an error.
129    pub fn read(state: &Path) -> Result<Self, String> {
130        let path = state.join(FILE);
131        let invalid = |message: String| format!("{}: {message}", crate::platform::path_text(&path));
132        let too_large =
133            || invalid("limits.json must be a regular file no larger than 64 KiB.".into());
134        let file = match crate::platform::open_regular(&path) {
135            Ok(file) => file,
136            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
137                return Ok(Self::default());
138            }
139            Err(error) if error.kind() == std::io::ErrorKind::InvalidInput => {
140                return Err(too_large());
141            }
142            Err(error) => return Err(invalid(error.to_string())),
143        };
144        let mut text = String::new();
145        std::io::Read::read_to_string(&mut std::io::Read::take(&file, MAX_BYTES + 1), &mut text)
146            .map_err(|e| invalid(e.to_string()))?;
147        if text.len() as u64 > MAX_BYTES {
148            return Err(too_large());
149        }
150        let declared: Declared = serde_json::from_str(&text).map_err(|error| {
151            let message = error.to_string();
152            // Backend validation previously followed decoding, so its diagnostics had no
153            // JSON location suffix. Preserve that contract when parsing names at the edge.
154            let message = if message.starts_with("backends:") || message.starts_with("backends.") {
155                message
156                    .strip_suffix(&format!(
157                        " at line {} column {}",
158                        error.line(),
159                        error.column()
160                    ))
161                    .unwrap_or(&message)
162            } else {
163                &message
164            };
165            invalid(message.to_owned())
166        })?;
167        let agent_sessions = declared
168            .agent_sessions
169            .map(DeclaredSessions::resolve)
170            .transpose()
171            .map_err(invalid)?
172            .unwrap_or_default();
173        Ok(Self {
174            backends: declared.backends,
175            agent_sessions,
176        })
177    }
178
179    /// The slot count of a parsed backend; `None` when it is unlimited.
180    pub fn limit(&self, backend: Backend) -> Option<u32> {
181        self.backends.get(&backend).copied()
182    }
183
184    pub fn backends(&self) -> &BTreeMap<Backend, u32> {
185        &self.backends
186    }
187
188    pub fn agent_sessions(&self) -> AgentSessions {
189        self.agent_sessions
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196    use std::fs;
197
198    fn read(text: &str) -> Result<Limits, String> {
199        let directory = crate::test_os::tempdir();
200        fs::write(directory.path().join(FILE), text).unwrap();
201        Limits::read(directory.path())
202    }
203
204    #[test]
205    fn agent_sessions_default_and_validate() {
206        assert_eq!(
207            read("{}").unwrap().agent_sessions(),
208            AgentSessions {
209                enabled: true,
210                max_bytes: 1_073_741_824,
211                target_bytes: 805_306_368,
212            }
213        );
214        assert_eq!(
215            read(r#"{"agentSessions":{"enabled":false}}"#)
216                .unwrap()
217                .agent_sessions(),
218            AgentSessions {
219                enabled: false,
220                max_bytes: 1_073_741_824,
221                target_bytes: 805_306_368,
222            }
223        );
224        assert_eq!(
225            read(r#"{"agentSessions":{"maxBytes":4000,"targetBytes":1000}}"#)
226                .unwrap()
227                .agent_sessions(),
228            AgentSessions {
229                enabled: true,
230                max_bytes: 4000,
231                target_bytes: 1000
232            }
233        );
234        // A small maxBytes alone collects down to three quarters of it.
235        assert_eq!(
236            read(r#"{"agentSessions":{"maxBytes":4000}}"#)
237                .unwrap()
238                .agent_sessions()
239                .target_bytes,
240            3000
241        );
242        // Quarter-block rounding stays divide-then-multiply, not floor(3*n/4).
243        for (max_bytes, target_bytes) in [(1, 0), (3, 0), (5, 3), (7, 3), (8, 6), (11, 6)] {
244            assert_eq!(
245                read(&format!(
246                    r#"{{"agentSessions":{{"maxBytes":{max_bytes}}}}}"#
247                ))
248                .unwrap()
249                .agent_sessions()
250                .target_bytes,
251                target_bytes
252            );
253        }
254        assert_eq!(
255            read(r#"{"agentSessions":{"maxBytes":18446744073709551615}}"#)
256                .unwrap()
257                .agent_sessions()
258                .target_bytes,
259            805306368
260        );
261        for (text, error) in [
262            (
263                r#"{"agentSessions":{"maxBytes":100,"targetBytes":100}}"#,
264                "agentSessions.targetBytes (100) must be lower than maxBytes (100).",
265            ),
266            (
267                r#"{"agentSessions":{"targetBytes":2147483648}}"#,
268                "must be lower than maxBytes (1073741824)",
269            ),
270            (
271                r#"{"agentSessions":{"maxBytes":0}}"#,
272                "agentSessions.maxBytes must be at least 1.",
273            ),
274            (r#"{"agentSessions":{"maxbytes":10}}"#, "unknown field"),
275            (r#"{"agentSessions":{"maxBytes":-1}}"#, "invalid value"),
276            (r#"{"agentSessions":{"enabled":"no"}}"#, "invalid type"),
277            (r#"{"agentSessions":true}"#, "invalid type"),
278        ] {
279            let message = read(text).unwrap_err();
280            assert!(message.contains(error), "{text}: {message}");
281            assert!(message.contains(FILE), "{message}");
282        }
283    }
284}