Skip to main content

a3s_code_core/session_review/
store.rs

1//! Bounded session-owned collection of durable review findings.
2
3use 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/// One session's durable review findings (sticky reply / science review).
12#[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    /// Insert or replace by `finding_id`. New findings must be `pending`.
53    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    /// Findings that must be injected into the next main-agent user prompt.
133    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    /// Pending findings for one registered scenario (host inject / UI filters).
141    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    /// Findings waiting for reviewer acceptance after the main agent addressed them.
151    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    /// Remap ownership when forking a session snapshot.
167    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        // Direct upsert of a brand-new non-pending finding is rejected.
247        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}