a3s_code_core/session_review/
store.rs1use super::{
4 validate_id, SessionReviewError, SessionReviewFindingV1, SessionReviewStatusV1,
5 SESSION_REVIEW_MAX_FINDINGS, SESSION_REVIEW_MAX_MESSAGE_BYTES,
6};
7use serde::{Deserialize, Serialize};
8
9pub const SESSION_REVIEW_STORE_SCHEMA_V1: &str = "a3s.code.session-review-store.v1";
10
11#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
13#[serde(rename_all = "camelCase", deny_unknown_fields)]
14pub struct SessionReviewStoreV1 {
15 pub schema: String,
16 pub session_id: String,
17 pub findings: Vec<SessionReviewFindingV1>,
18}
19
20impl SessionReviewStoreV1 {
21 pub fn empty(session_id: impl Into<String>) -> Result<Self, SessionReviewError> {
22 let store = Self {
23 schema: SESSION_REVIEW_STORE_SCHEMA_V1.to_owned(),
24 session_id: session_id.into(),
25 findings: Vec::new(),
26 };
27 store.validate()?;
28 Ok(store)
29 }
30
31 pub fn validate(&self) -> Result<(), SessionReviewError> {
32 if self.schema != SESSION_REVIEW_STORE_SCHEMA_V1 {
33 return Err(SessionReviewError::UnsupportedSchema);
34 }
35 validate_id("sessionId", &self.session_id)?;
36 if self.findings.len() > SESSION_REVIEW_MAX_FINDINGS {
37 return Err(SessionReviewError::InvalidField("findings"));
38 }
39 let mut seen = std::collections::BTreeSet::new();
40 for finding in &self.findings {
41 finding.validate()?;
42 if finding.session_id != self.session_id {
43 return Err(SessionReviewError::InvalidField("finding.sessionId"));
44 }
45 if !seen.insert(finding.finding_id.clone()) {
46 return Err(SessionReviewError::InvalidField("findingId"));
47 }
48 }
49 Ok(())
50 }
51
52 pub fn upsert(&mut self, finding: SessionReviewFindingV1) -> Result<(), SessionReviewError> {
54 finding.validate()?;
55 if finding.session_id != self.session_id {
56 return Err(SessionReviewError::InvalidField("finding.sessionId"));
57 }
58 if let Some(existing) = self
59 .findings
60 .iter()
61 .position(|entry| entry.finding_id == finding.finding_id)
62 {
63 self.findings[existing] = finding;
64 } else {
65 if self.findings.len() >= SESSION_REVIEW_MAX_FINDINGS {
66 return Err(SessionReviewError::InvalidField("findings"));
67 }
68 if !matches!(finding.status, SessionReviewStatusV1::Pending) {
69 return Err(SessionReviewError::InvalidField("finding.status"));
70 }
71 self.findings.push(finding);
72 self.findings
73 .sort_unstable_by(|left, right| left.finding_id.cmp(&right.finding_id));
74 }
75 self.validate()
76 }
77
78 pub fn get_mut(
79 &mut self,
80 finding_id: &str,
81 ) -> Result<&mut SessionReviewFindingV1, SessionReviewError> {
82 self.findings
83 .iter_mut()
84 .find(|finding| finding.finding_id == finding_id)
85 .ok_or_else(|| SessionReviewError::FindingNotFound(finding_id.to_owned()))
86 }
87
88 pub fn mark_addressed(
89 &mut self,
90 finding_id: &str,
91 run_id: impl Into<String>,
92 at_ms: u64,
93 ) -> Result<(), SessionReviewError> {
94 self.get_mut(finding_id)?.mark_addressed(run_id, at_ms)?;
95 self.validate()
96 }
97
98 pub fn set_main_agent_reply(
99 &mut self,
100 finding_id: &str,
101 reply: impl Into<String>,
102 ) -> Result<(), SessionReviewError> {
103 self.get_mut(finding_id)?.set_main_agent_reply(reply)?;
104 self.validate()
105 }
106
107 pub fn accept(
108 &mut self,
109 finding_id: &str,
110 review_id: impl Into<String>,
111 at_ms: u64,
112 ) -> Result<(), SessionReviewError> {
113 self.get_mut(finding_id)?.accept(review_id, at_ms)?;
114 self.validate()
115 }
116
117 pub fn reopen(
118 &mut self,
119 finding_id: &str,
120 reason: impl Into<String>,
121 at_ms: u64,
122 ) -> Result<(), SessionReviewError> {
123 self.get_mut(finding_id)?.reopen(reason, at_ms)?;
124 self.validate()
125 }
126
127 pub fn waive(&mut self, finding_id: &str, at_ms: u64) -> Result<(), SessionReviewError> {
128 self.get_mut(finding_id)?.waive(at_ms)?;
129 self.validate()
130 }
131
132 pub fn pending_for_injection(&self) -> Vec<&SessionReviewFindingV1> {
134 self.findings
135 .iter()
136 .filter(|finding| finding.status.injects_into_main_prompt())
137 .collect()
138 }
139
140 pub fn pending_for_scenario(&self, scenario_id: &str) -> Vec<&SessionReviewFindingV1> {
142 self.findings
143 .iter()
144 .filter(|finding| {
145 finding.scenario_id == scenario_id && finding.status.injects_into_main_prompt()
146 })
147 .collect()
148 }
149
150 pub fn awaiting_acceptance(&self) -> Vec<&SessionReviewFindingV1> {
152 self.findings
153 .iter()
154 .filter(|finding| matches!(finding.status, SessionReviewStatusV1::Addressed))
155 .collect()
156 }
157
158 pub fn pending_count(&self) -> usize {
159 self.pending_for_injection().len()
160 }
161
162 pub fn addressed_count(&self) -> usize {
163 self.awaiting_acceptance().len()
164 }
165
166 pub fn rebind_session(
168 &mut self,
169 session_id: impl Into<String>,
170 ) -> Result<(), SessionReviewError> {
171 let session_id = session_id.into();
172 validate_id("sessionId", &session_id)?;
173 self.session_id = session_id.clone();
174 for finding in &mut self.findings {
175 finding.session_id = session_id.clone();
176 }
177 self.validate()
178 }
179
180 pub fn from_slice(bytes: &[u8]) -> Result<Self, SessionReviewError> {
181 let store: Self = super::decode_json_slice(bytes)?;
182 store.validate()?;
183 Ok(store)
184 }
185
186 pub fn to_vec(&self) -> Result<Vec<u8>, SessionReviewError> {
187 self.validate()?;
188 let bytes = super::encode_json(self)?;
189 if bytes.len() > SESSION_REVIEW_MAX_MESSAGE_BYTES {
190 return Err(SessionReviewError::Encoding);
191 }
192 Ok(bytes)
193 }
194}
195
196#[cfg(test)]
197mod tests {
198 use super::*;
199 use crate::session_review::{SessionReviewAnchorV1, SessionReviewSeverityV1};
200
201 fn pending_finding(id: &str) -> SessionReviewFindingV1 {
202 SessionReviewFindingV1::new_transcript(
203 id,
204 "session-1",
205 SessionReviewSeverityV1::Warning,
206 "numeric",
207 "claim",
208 "evidence",
209 SessionReviewAnchorV1::new("turn-1").unwrap(),
210 "review-1",
211 1_000,
212 )
213 .unwrap()
214 }
215
216 #[test]
217 fn upsert_and_injection_queues() {
218 let mut store = SessionReviewStoreV1::empty("session-1").unwrap();
219 store.upsert(pending_finding("f-1")).unwrap();
220 store.upsert(pending_finding("f-2")).unwrap();
221 assert_eq!(store.pending_count(), 2);
222 store.mark_addressed("f-1", "run-9", 1_100).unwrap();
223 assert_eq!(store.pending_count(), 1);
224 assert_eq!(store.addressed_count(), 1);
225 assert_eq!(store.pending_for_injection()[0].finding_id, "f-2");
226 store.accept("f-1", "review-2", 1_200).unwrap();
227 assert_eq!(store.addressed_count(), 0);
228 assert!(!store
229 .pending_for_injection()
230 .iter()
231 .any(|finding| finding.finding_id == "f-1"));
232 }
233
234 #[test]
235 fn reject_foreign_session_and_duplicate_new_non_pending() {
236 let mut store = SessionReviewStoreV1::empty("session-1").unwrap();
237 let mut foreign = pending_finding("f-1");
238 foreign.session_id = "other".into();
239 assert_eq!(
240 store.upsert(foreign),
241 Err(SessionReviewError::InvalidField("finding.sessionId"))
242 );
243
244 let mut addressed = pending_finding("f-2");
245 addressed.mark_addressed("run-1", 1_100).unwrap();
246 let mut store2 = SessionReviewStoreV1::empty("session-1").unwrap();
248 assert_eq!(
249 store2.upsert(addressed),
250 Err(SessionReviewError::InvalidField("finding.status"))
251 );
252 }
253
254 #[test]
255 fn rebind_session_updates_all_findings() {
256 let mut store = SessionReviewStoreV1::empty("session-1").unwrap();
257 store.upsert(pending_finding("f-1")).unwrap();
258 store.rebind_session("session-fork").unwrap();
259 assert_eq!(store.session_id, "session-fork");
260 assert_eq!(store.findings[0].session_id, "session-fork");
261 }
262
263 #[test]
264 fn missing_finding_errors() {
265 let mut store = SessionReviewStoreV1::empty("session-1").unwrap();
266 assert!(matches!(
267 store.waive("missing", 1),
268 Err(SessionReviewError::FindingNotFound(_))
269 ));
270 }
271
272 #[test]
273 fn store_persists_main_agent_reply_until_reopen() {
274 let mut store = SessionReviewStoreV1::empty("session-1").unwrap();
275 store.upsert(pending_finding("f-1")).unwrap();
276 store.mark_addressed("f-1", "run-9", 1_100).unwrap();
277 store
278 .set_main_agent_reply("f-1", " Fixed via tool output ")
279 .unwrap();
280 assert_eq!(
281 store.findings[0].main_agent_reply.as_deref(),
282 Some("Fixed via tool output")
283 );
284 assert_eq!(
285 store.awaiting_acceptance()[0].main_agent_reply.as_deref(),
286 Some("Fixed via tool output")
287 );
288 assert!(matches!(
289 store.set_main_agent_reply("missing", "x"),
290 Err(SessionReviewError::FindingNotFound(_))
291 ));
292 store
293 .reopen("f-1", "evidence still incomplete", 1_200)
294 .unwrap();
295 assert!(store.findings[0].main_agent_reply.is_none());
296 }
297}