1use std::collections::BTreeMap;
4
5use chrono::{DateTime, TimeZone, Utc};
6
7use super::{
8 Result,
9 dictionary::{AgentKey, PrincipalKey},
10 invalid,
11 io::Reader,
12 limits::{
13 MAX_COMPACT_STATE_COUNT, MIN_AGENT_BYTES, MIN_EXTRA_HEADER_BYTES, MIN_LINEAGE_BYTES,
14 MIN_PRINCIPAL_BYTES, MIN_STATE_COLUMN_BYTES, MIN_STATE_PARENT_BYTES,
15 MIN_VERIFICATION_CUSTOM_BYTES, admit_count,
16 },
17 state::{STATE_MAGIC, STATE_MAGIC_V1},
18};
19use crate::object::{
20 Agent, Attribution, ChangeId, ChangeLineage, ChangeLineageKind, ContentHash, Principal, State,
21 StateId, Status, Verification,
22};
23
24pub fn decode_state_frame(bytes: &[u8]) -> Result<Vec<State>> {
26 let magic = state_frame_magic(bytes)?;
27 let mut input = Reader::verified(bytes, magic)?;
28 let count = input.get_count_at_most("state frame", 1, MAX_COMPACT_STATE_COUNT)?;
29 let (principals, agents) = decode_dictionaries(&mut input, magic == STATE_MAGIC)?;
30 admit_count(
31 "state frame",
32 count,
33 input.remaining(),
34 MIN_STATE_COLUMN_BYTES,
35 MAX_COMPACT_STATE_COUNT,
36 )?;
37 let mut states = (0..count).map(|_| blank_state()).collect::<Vec<_>>();
38 decode_structure(&mut input, &mut states)?;
39 decode_attribution(&mut input, &mut states, &principals, &agents)?;
40 decode_intent_and_verification(&mut input, &mut states)?;
41 decode_timestamps(&mut input, &mut states)?;
42 decode_fidelity(&mut input, &mut states, &principals)?;
43 decode_lineage(&mut input, &mut states)?;
44 input.finish()?;
45 for state in &mut states {
46 state.state_id = state.id();
47 }
48 Ok(states)
49}
50
51fn decode_structure(input: &mut Reader<'_>, states: &mut [State]) -> Result<()> {
52 for state in &mut *states {
53 state.change_id = ChangeId::from_bytes(input.get_fixed()?);
54 }
55 for state in &mut *states {
56 state.tree = ContentHash::from_bytes(input.get_fixed()?);
57 }
58 for state in states {
59 let count = input.get_count("state parent", MIN_STATE_PARENT_BYTES)?;
60 state.parents = (0..count)
61 .map(|_| Ok(StateId::from_bytes(input.get_fixed()?)))
62 .collect::<Result<Vec<_>>>()?;
63 }
64 Ok(())
65}
66
67fn decode_attribution(
68 input: &mut Reader<'_>,
69 states: &mut [State],
70 principals: &[Principal],
71 agents: &[Agent],
72) -> Result<()> {
73 for state in &mut *states {
74 state.attribution.principal = principal_at(principals, input.get_u64()?)?;
75 }
76 for state in states {
77 state.attribution.agent = optional_at(agents, input.get_u64()?)?;
78 }
79 Ok(())
80}
81
82fn decode_intent_and_verification(input: &mut Reader<'_>, states: &mut [State]) -> Result<()> {
83 for state in &mut *states {
84 state.intent = input
85 .get_optional_bytes()?
86 .map(String::from_utf8)
87 .transpose()
88 .map_err(|_| invalid("state intent is not UTF-8"))?;
89 }
90 for state in &mut *states {
91 state.confidence = get_optional_f32(input)?;
92 }
93 for state in &mut *states {
94 state.verification = decode_verification(input)?;
95 }
96 for state in states {
97 state.status =
98 Status::from_byte(input.get_u8()?).ok_or_else(|| invalid("invalid state status"))?;
99 }
100 Ok(())
101}
102
103fn decode_timestamps(input: &mut Reader<'_>, states: &mut [State]) -> Result<()> {
104 let mut previous = 0i64;
105 for (index, state) in states.iter_mut().enumerate() {
106 let encoded = input.get_i64()?;
107 let seconds = if index == 0 {
108 encoded
109 } else {
110 previous
111 .checked_add(encoded)
112 .ok_or_else(|| invalid("created timestamp delta overflow"))?
113 };
114 state.created_at = timestamp(seconds, get_u32(input, "created timestamp nanos")?)?;
115 previous = seconds;
116 }
117 for state in &mut *states {
118 state.authored_at = match input.get_u8()? {
119 0 => None,
120 1 => {
121 let seconds = state
122 .created_at
123 .timestamp()
124 .checked_add(input.get_i64()?)
125 .ok_or_else(|| invalid("authored timestamp delta overflow"))?;
126 Some(timestamp(
127 seconds,
128 get_u32(input, "authored timestamp nanos")?,
129 )?)
130 }
131 value => return Err(invalid(format!("invalid authored timestamp tag {value}"))),
132 };
133 }
134 for state in states {
135 state.authored_tz_offset = get_i32(input, "author timezone")?;
136 state.committer_tz_offset = get_i32(input, "committer timezone")?;
137 }
138 Ok(())
139}
140
141fn decode_fidelity(
142 input: &mut Reader<'_>,
143 states: &mut [State],
144 principals: &[Principal],
145) -> Result<()> {
146 for state in &mut *states {
147 state.provenance = match input.get_u8()? {
148 0 => None,
149 1 => Some(ContentHash::from_bytes(input.get_fixed()?)),
150 value => return Err(invalid(format!("invalid provenance option tag {value}"))),
151 };
152 }
153 for state in &mut *states {
154 state.committer = optional_at(principals, input.get_u64()?)?;
155 }
156 for state in &mut *states {
157 state.raw_message = input.get_optional_bytes()?;
158 }
159 for state in &mut *states {
160 let count = input.get_count("extra header", MIN_EXTRA_HEADER_BYTES)?;
161 state.extra_headers = (0..count)
162 .map(|_| Ok((input.get_bytes()?, input.get_bytes()?)))
163 .collect::<Result<Vec<_>>>()?;
164 }
165 for state in states {
166 state.git_lossy = input.get_bool()?;
167 }
168 Ok(())
169}
170
171fn decode_lineage(input: &mut Reader<'_>, states: &mut [State]) -> Result<()> {
172 for state in states {
173 let count = input.get_count("state lineage", MIN_LINEAGE_BYTES)?;
174 state.lineage = (0..count)
175 .map(|_| {
176 Ok(ChangeLineage {
177 kind: decode_lineage_kind(input.get_u8()?)?,
178 source_change: ChangeId::from_bytes(input.get_fixed()?),
179 source_state: StateId::from_bytes(input.get_fixed()?),
180 })
181 })
182 .collect::<Result<Vec<_>>>()?;
183 }
184 Ok(())
185}
186
187fn decode_verification(input: &mut Reader<'_>) -> Result<Option<Verification>> {
188 match input.get_u8()? {
189 0 => Ok(None),
190 1 => {
191 let tests_passed = get_optional_bool(input)?;
192 let tests_failed = get_optional_u32(input)?;
193 let coverage_pct = get_optional_f32(input)?;
194 let coverage_delta = get_optional_f32(input)?;
195 let lint_warnings = get_optional_u32(input)?;
196 let count = input.get_count("verification custom", MIN_VERIFICATION_CUSTOM_BYTES)?;
197 let mut custom = BTreeMap::new();
198 for _ in 0..count {
199 let key = String::from_utf8(input.get_bytes()?)
200 .map_err(|_| invalid("verification key is not UTF-8"))?;
201 custom.insert(key, rmp_serde::from_slice(&input.get_bytes()?)?);
202 }
203 Ok(Some(Verification {
204 tests_passed,
205 tests_failed,
206 coverage_pct,
207 coverage_delta,
208 lint_warnings,
209 custom,
210 }))
211 }
212 value => Err(invalid(format!("invalid verification option tag {value}"))),
213 }
214}
215
216fn state_frame_magic(bytes: &[u8]) -> Result<&'static [u8; 4]> {
217 if bytes.starts_with(STATE_MAGIC) {
218 Ok(STATE_MAGIC)
219 } else if bytes.starts_with(STATE_MAGIC_V1) {
220 Ok(STATE_MAGIC_V1)
221 } else {
222 Err(invalid("frame magic does not match its object kind"))
223 }
224}
225
226fn decode_dictionaries(
227 input: &mut Reader<'_>,
228 include_cursor_fields: bool,
229) -> Result<(Vec<Principal>, Vec<Agent>)> {
230 let principals = (0..input.get_count("principal dictionary", MIN_PRINCIPAL_BYTES)?)
231 .map(|_| principal_from_key(PrincipalKey(input.get_bytes()?, input.get_bytes()?)))
232 .collect::<Result<Vec<_>>>()?;
233 let min_agent_bytes = if include_cursor_fields {
234 MIN_AGENT_BYTES
235 } else {
236 MIN_AGENT_BYTES.saturating_sub(2)
237 };
238 let agents = (0..input.get_count("agent dictionary", min_agent_bytes)?)
239 .map(|_| {
240 agent_from_key(AgentKey {
241 provider: input.get_bytes()?,
242 model: input.get_bytes()?,
243 session_id: input.get_optional_bytes()?,
244 segment_id: input.get_optional_bytes()?,
245 policy_id: input.get_optional_bytes()?,
246 thought_level: if include_cursor_fields {
247 input.get_optional_bytes()?
248 } else {
249 None
250 },
251 parent: if include_cursor_fields {
252 input.get_optional_bytes()?
253 } else {
254 None
255 },
256 })
257 })
258 .collect::<Result<Vec<_>>>()?;
259 Ok((principals, agents))
260}
261
262fn principal_from_key(value: PrincipalKey) -> Result<Principal> {
263 Ok(Principal {
264 name: value.0,
265 email: value.1,
266 })
267}
268
269fn agent_from_key(value: AgentKey) -> Result<Agent> {
270 Ok(Agent {
271 provider: String::from_utf8(value.provider)
272 .map_err(|_| invalid("agent provider is not UTF-8"))?,
273 model: String::from_utf8(value.model).map_err(|_| invalid("agent model is not UTF-8"))?,
274 session_id: optional_string(value.session_id, "agent session id")?,
275 segment_id: optional_string(value.segment_id, "agent segment id")?,
276 policy_id: optional_string(value.policy_id, "agent policy id")?,
277 thought_level: optional_string(value.thought_level, "agent thought_level")?,
278 parent: optional_string(value.parent, "agent parent")?,
279 })
280}
281
282fn optional_string(value: Option<Vec<u8>>, field: &str) -> Result<Option<String>> {
283 value
284 .map(String::from_utf8)
285 .transpose()
286 .map_err(|_| invalid(format!("{field} is not UTF-8")))
287}
288
289fn principal_at(values: &[Principal], index: u64) -> Result<Principal> {
290 values
291 .get(index_usize(index)?)
292 .cloned()
293 .ok_or_else(|| invalid("principal dictionary index is out of range"))
294}
295
296fn optional_at<T: Clone>(values: &[T], encoded: u64) -> Result<Option<T>> {
297 if encoded == 0 {
298 return Ok(None);
299 }
300 Ok(Some(
301 values
302 .get(index_usize(encoded - 1)?)
303 .cloned()
304 .ok_or_else(|| invalid("dictionary index is out of range"))?,
305 ))
306}
307
308fn index_usize(value: u64) -> Result<usize> {
309 usize::try_from(value).map_err(|_| invalid("dictionary index exceeds platform limits"))
310}
311
312fn get_optional_bool(input: &mut Reader<'_>) -> Result<Option<bool>> {
313 match input.get_u8()? {
314 0 => Ok(None),
315 1 => Ok(Some(false)),
316 2 => Ok(Some(true)),
317 value => Err(invalid(format!("invalid optional boolean {value}"))),
318 }
319}
320
321fn get_optional_u32(input: &mut Reader<'_>) -> Result<Option<u32>> {
322 match input.get_u64()? {
323 0 => Ok(None),
324 value => Ok(Some(
325 u32::try_from(value - 1).map_err(|_| invalid("optional u32 exceeds its range"))?,
326 )),
327 }
328}
329
330fn get_optional_f32(input: &mut Reader<'_>) -> Result<Option<f32>> {
331 match input.get_u8()? {
332 0 => Ok(None),
333 1 => Ok(Some(f32::from_le_bytes(input.get_fixed()?))),
334 value => Err(invalid(format!("invalid optional f32 tag {value}"))),
335 }
336}
337
338fn get_u32(input: &mut Reader<'_>, field: &str) -> Result<u32> {
339 u32::try_from(input.get_u64()?).map_err(|_| invalid(format!("{field} exceeds u32")))
340}
341
342fn get_i32(input: &mut Reader<'_>, field: &str) -> Result<i32> {
343 i32::try_from(input.get_i64()?).map_err(|_| invalid(format!("{field} exceeds i32")))
344}
345
346fn timestamp(seconds: i64, nanos: u32) -> Result<DateTime<Utc>> {
347 Utc.timestamp_opt(seconds, nanos)
348 .single()
349 .ok_or_else(|| invalid("compact timestamp is out of range"))
350}
351
352fn decode_lineage_kind(value: u8) -> Result<ChangeLineageKind> {
353 match value {
354 1 => Ok(ChangeLineageKind::CherryPick),
355 2 => Ok(ChangeLineageKind::Collapse),
356 3 => Ok(ChangeLineageKind::Revert),
357 4 => Ok(ChangeLineageKind::GitProjection),
358 value => Err(invalid(format!("invalid lineage kind {value}"))),
359 }
360}
361
362fn blank_state() -> State {
363 State {
364 state_id: StateId::default(),
365 change_id: ChangeId::from_bytes([0; 16]),
366 tree: ContentHash::from_bytes([0; 32]),
367 parents: Vec::new(),
368 attribution: Attribution::human(Principal::new("", "")),
369 intent: None,
370 confidence: None,
371 created_at: DateTime::UNIX_EPOCH,
372 verification: None,
373 status: Status::Draft,
374 provenance: None,
375 authored_at: None,
376 committer: None,
377 authored_tz_offset: 0,
378 committer_tz_offset: 0,
379 raw_message: None,
380 git_lossy: false,
381 extra_headers: Vec::new(),
382 lineage: Vec::new(),
383 }
384}