supercode_interchange/session/

residue.rs

//! The portable residue as the complement of a projection (`docs/plans/portable-residue.md`).
//!
//! A translation keeps what the target format can say; the residue is the rest of the source.
//! A source file is cut into **segments**, one per message group (a message and the tool results
//! answering it), each holding the verbatim source records from the record that created the group
//! up to the next group's creating record. Segments are keyed by a hash chain over the groups'
//! canonical form, which every format reproduces, so a session read back from any format finds its
//! source's segments and restores that prefix byte for byte; messages past the restored prefix are
//! written by the format's own splice writer.

use super::*;

/// Loader-internal: the index in `raw` of the record that created a message, until the loader
/// moves it into [`SessionMeta::message_records`] with [`take_message_records`].
const RECORD_KEY: &str = "__supercode_record";

/// Stamp the messages created since the last stamp with `record`. They are always a suffix: a
/// loader appends, or rebuilds the list outright (Codex's `compacted` replacement history).
pub(super) fn stamp_message_records(messages: &mut [ChatMessage], record: usize) {
    for message in messages.iter_mut().rev() {
        if message.metadata.contains_key(RECORD_KEY) {
            break;
        }
        message
            .metadata
            .insert(RECORD_KEY.to_string(), record.to_string());
    }
}

/// Remove the stamps and return them in message order; a message without one (added after
/// parsing, e.g. a synthesized tool result) has `None`.
pub(super) fn take_message_records(messages: &mut [ChatMessage]) -> Vec<Option<usize>> {
    messages
        .iter_mut()
        .map(|message| {
            message
                .metadata
                .remove(RECORD_KEY)
                .and_then(|record| record.parse().ok())
        })
        .collect()
}

/// One stored segment: the verbatim source lines of a message group.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResidueSegment {
    /// The verbatim source lines, in order.
    pub lines: Vec<String>,
    /// The segment ends its source file (whose trailing newline it records).
    pub end_of_source: bool,
    /// Whether that source file ended with a newline.
    pub trailing_newline: bool,
    /// BLAKE3 over the lines, checked before a segment is used: a segment altered or damaged in
    /// the store is treated as missing.
    #[serde(default)]
    pub digest: String,
    /// The source session the lines were cut from. Sessions can share a conversation prefix
    /// (sub-agents opened with one prompt, forks), so a key can hold segments of several sources,
    /// and a restoration takes all of its segments from one.
    #[serde(default)]
    pub source: String,
    /// That source's session id, when it has one.
    #[serde(default)]
    pub source_session: Option<String>,
}

impl ResidueSegment {
    fn new(
        lines: Vec<String>,
        end_of_source: bool,
        trailing_newline: bool,
        source: &str,
        source_session: Option<&str>,
    ) -> Self {
        let digest = segment_digest(&lines);
        Self {
            lines,
            end_of_source,
            trailing_newline,
            digest,
            source: source.to_string(),
            source_session: source_session.map(str::to_string),
        }
    }

    /// Whether the lines are the ones the segment was cut with.
    pub fn verified(&self) -> bool {
        self.digest == segment_digest(&self.lines)
    }
}

fn segment_digest(lines: &[String]) -> String {
    let mut hasher = blake3::Hasher::new();
    for line in lines {
        hasher.update(&(line.len() as u64).to_le_bytes());
        hasher.update(line.as_bytes());
    }
    hasher.finalize().to_hex().to_string()
}

/// A segment ready to store: its chain key, its source format and its lines.
#[derive(Debug, Clone)]
pub struct KeyedSegment {
    /// The chain key of the segment's message group.
    pub key: String,
    /// The format the lines are in.
    pub format: SessionFormat,
    /// The lines themselves.
    pub segment: ResidueSegment,
}

/// The message groups of `messages`: each non-tool message opens a group; tool results join the
/// group of the calls they answer.
fn message_groups(messages: &[ChatMessage]) -> Vec<std::ops::Range<usize>> {
    let mut groups: Vec<std::ops::Range<usize>> = Vec::new();
    for (index, message) in messages.iter().enumerate() {
        match groups.last_mut() {
            Some(group) if message.role == Role::Tool => group.end = index + 1,
            _ => groups.push(index..index + 1),
        }
    }
    groups
}

fn message_text(message: &ChatMessage) -> String {
    match &message.content_parts {
        Some(parts) => parts
            .iter()
            .filter(|part| part.get("type").and_then(Value::as_str) == Some("text"))
            .filter_map(|part| part.get("text").and_then(Value::as_str))
            .collect(),
        None => message.content.clone().unwrap_or_default(),
    }
}

fn normalized_arguments(arguments: &str) -> Value {
    serde_json::from_str(arguments).unwrap_or_else(|_| Value::String(arguments.to_string()))
}

/// The canonical form of a group: what every format carries of it. Tool results are ordered by
/// the calls they answer, since formats store parallel results in call or completion order.
fn canonical_group(messages: &[ChatMessage]) -> Value {
    let head = &messages[0];
    let calls = head.tool_calls();
    let mut results: Vec<&ChatMessage> = messages[1..].iter().collect();
    results.sort_by_key(|result| {
        calls
            .iter()
            .position(|call| Some(call.id.as_str()) == result.tool_call_id.as_deref())
            .unwrap_or(usize::MAX)
    });
    serde_json::json!([
        serde_json::to_value(head.role).unwrap_or(Value::Null),
        message_text(head),
        calls
            .iter()
            .map(|call| serde_json::json!([
                call.function.name,
                normalized_arguments(&call.function.arguments)
            ]))
            .collect::<Vec<_>>(),
        results
            .iter()
            .map(|result| message_text(result))
            .collect::<Vec<_>>(),
    ])
}

/// The key before any group: an empty conversation's only key. Sessions with no messages share
/// it, so only their session id tells them apart.
fn empty_conversation_key() -> String {
    blake3::hash(b"supercode-residue-v1").to_hex().to_string()
}

/// The chain key of every group: `key(g) = H(key(g-1), canonical(g))`.
pub fn residue_chain_keys(messages: &[ChatMessage]) -> Vec<(std::ops::Range<usize>, String)> {
    let mut previous = empty_conversation_key();
    message_groups(messages)
        .into_iter()
        .map(|group| {
            let mut hasher = blake3::Hasher::new();
            hasher.update(previous.as_bytes());
            hasher.update(
                canonical_group(&messages[group.clone()])
                    .to_string()
                    .as_bytes(),
            );
            previous = hasher.finalize().to_hex().to_string();
            (group, previous.clone())
        })
        .collect()
}

impl Session {
    /// This session's segments, when its `raw` is its verbatim source and its loader recorded
    /// which record created each message. Segments stop at the first group whose creating record
    /// is unknown; the first segment also holds every record before the first message.
    pub fn residue_segments(&self) -> Vec<KeyedSegment> {
        let Some(format) = SessionFormat::for_source(self.meta.source) else {
            return Vec::new();
        };
        if !self.raw_is_verbatim || self.meta.message_records.len() != self.messages.len() {
            return Vec::new();
        }
        let keyed = residue_chain_keys(&self.messages);
        // The source's identity: its format, its id and its first record.
        let source = {
            let mut hasher = blake3::Hasher::new();
            hasher.update(format!("{format:?}").as_bytes());
            hasher.update(self.meta.session_id.as_deref().unwrap_or("").as_bytes());
            hasher.update(
                self.raw
                    .first()
                    .map(String::as_str)
                    .unwrap_or("")
                    .as_bytes(),
            );
            hasher.finalize().to_hex().to_string()
        };
        if keyed.is_empty() {
            if self.raw.is_empty() {
                return Vec::new();
            }
            return vec![KeyedSegment {
                key: empty_conversation_key(),
                format,
                segment: ResidueSegment::new(
                    self.raw.clone(),
                    true,
                    self.raw_trailing_newline,
                    &source,
                    self.meta.session_id.as_deref(),
                ),
            }];
        }
        let mut starts = Vec::new();
        for (group, _) in &keyed {
            match self.meta.message_records[group.start] {
                Some(record) => starts.push(record),
                None => break,
            }
        }
        let mut segments = Vec::new();
        for (index, (_, key)) in keyed.iter().take(starts.len()).enumerate() {
            let from = if index == 0 { 0 } else { starts[index] };
            let end_of_source = index + 1 == keyed.len();
            let to = if end_of_source {
                self.raw.len()
            } else {
                starts.get(index + 1).copied().unwrap_or(self.raw.len())
            };
            if from > to || to > self.raw.len() {
                break;
            }
            segments.push(KeyedSegment {
                key: key.clone(),
                format,
                segment: ResidueSegment::new(
                    self.raw[from..to].to_vec(),
                    end_of_source,
                    self.raw_trailing_newline,
                    &source,
                    self.meta.session_id.as_deref(),
                ),
            });
        }
        segments
    }

    /// Write this session as `format`, restoring the longest prefix whose segments `lookup` holds
    /// for that format verbatim and writing the rest with the format's splice writer. `None` when
    /// no prefix is found; the caller writes the session as usual.
    pub fn restore_residue(
        &self,
        format: SessionFormat,
        lookup: impl Fn(SessionFormat, &str) -> Vec<ResidueSegment>,
    ) -> Result<Option<(String, ResidueRestore)>> {
        let keyed = residue_chain_keys(&self.messages);
        if keyed.is_empty() {
            return Ok(self.restore_empty_conversation(format, &lookup));
        }
        let Some(source) = self.residue_source(format, &lookup, &keyed) else {
            return Ok(None);
        };
        let lookup = |format: SessionFormat, key: &str| {
            lookup(format, key)
                .into_iter()
                .find(|segment| segment.source == source && segment.verified())
        };
        let mut limit = keyed.len();
        // A segment that fails the soundness check costs only itself and what follows: retry
        // with the prefix before the first group that did not read back.
        loop {
            match self.restore_prefix(format, &lookup, &keyed, limit)? {
                Restored::Sound(text, report) => return Ok(Some((text, report))),
                Restored::Nothing => return Ok(None),
                Restored::Unsound {
                    first_mismatch,
                    restored_groups,
                } => {
                    let shorter = first_mismatch.min(restored_groups.saturating_sub(1));
                    if shorter == 0 || shorter >= limit {
                        return Ok(None);
                    }
                    limit = shorter;
                }
            }
        }
    }

    /// The one source session to restore from: never another session with a different recorded
    /// id; the one whose segments match the longest prefix of this conversation, then the one with
    /// this session's id. `None` when no source matches, or when sources tie with different lines
    /// — which session this is cannot be told apart.
    fn residue_source(
        &self,
        format: SessionFormat,
        lookup: &impl Fn(SessionFormat, &str) -> Vec<ResidueSegment>,
        keyed: &[(std::ops::Range<usize>, String)],
    ) -> Option<String> {
        let first = keyed.first()?;
        let own = self.meta.session_id.as_deref();
        // An identical conversation from another session is not this session: its lines carry
        // that session's id and workspace. Only an id the session's own records carry tells
        // sessions apart; Grok's is its directory, which a materialization mints anew.
        let own_recorded = own.filter(|_| self.meta.source != SessionSource::Grok);
        let mut best: Vec<(String, Option<String>, usize, Vec<String>)> = Vec::new();
        let candidates: Vec<ResidueSegment> = lookup(format, &first.1)
            .into_iter()
            .filter(ResidueSegment::verified)
            .filter(|candidate| {
                own_recorded.is_none()
                    || candidate.source_session.is_none()
                    || candidate.source_session.as_deref() == own_recorded
            })
            .collect();
        for candidate in &candidates {
            let mut matched = 0;
            let mut lines = Vec::new();
            for (_, key) in keyed {
                let Some(segment) = lookup(format, key)
                    .into_iter()
                    .find(|segment| segment.source == candidate.source && segment.verified())
                else {
                    break;
                };
                matched += 1;
                let end = segment.end_of_source;
                lines.extend(segment.lines);
                if end {
                    break;
                }
            }
            best.push((
                candidate.source.clone(),
                candidate.source_session.clone(),
                matched,
                lines,
            ));
        }
        let longest = best.iter().map(|(_, _, matched, _)| *matched).max()?;
        best.retain(|(_, _, matched, _)| *matched == longest);
        if best.len() > 1 {
            if own.is_some()
                && best
                    .iter()
                    .any(|(_, session, _, _)| session.as_deref() == own)
            {
                best.retain(|(_, session, _, _)| session.as_deref() == own);
            }
        }
        let (source, _, _, lines) = best.first()?;
        best.iter()
            .all(|(_, _, _, other)| other == lines)
            .then(|| source.clone())
    }

    /// A session with no messages: restored only when its id, or a single stored candidate,
    /// names which empty session it is.
    fn restore_empty_conversation(
        &self,
        format: SessionFormat,
        lookup: &impl Fn(SessionFormat, &str) -> Vec<ResidueSegment>,
    ) -> Option<(String, ResidueRestore)> {
        let mut candidates: Vec<ResidueSegment> = lookup(format, &empty_conversation_key())
            .into_iter()
            .filter(ResidueSegment::verified)
            .collect();
        let own = self.meta.session_id.as_deref();
        if candidates.len() > 1 && own.is_some() {
            candidates.retain(|segment| segment.source_session.as_deref() == own);
        }
        let [segment] = candidates.as_slice() else {
            return None;
        };
        let mut text = segment.lines.join("\n");
        if segment.trailing_newline {
            text.push('\n');
        }
        let written = Session::load_str(&text, format).ok()?;
        written.messages.is_empty().then_some((
            text,
            ResidueRestore {
                restored_messages: 0,
                rendered_messages: 0,
            },
        ))
    }

    fn restore_prefix(
        &self,
        format: SessionFormat,
        lookup: &impl Fn(SessionFormat, &str) -> Option<ResidueSegment>,
        keyed: &[(std::ops::Range<usize>, String)],
        limit: usize,
    ) -> Result<Restored> {
        let mut lines = Vec::new();
        let mut restored_groups = 0;
        let mut restored_messages = 0;
        let mut whole_source = None;
        for (group, key) in keyed.iter().take(limit) {
            let Some(segment) = lookup(format, key) else {
                break;
            };
            lines.extend(segment.lines);
            restored_groups += 1;
            restored_messages = group.end;
            if segment.end_of_source {
                whole_source = Some(segment.trailing_newline);
                break;
            }
        }
        if restored_groups == 0 {
            return Ok(Restored::Nothing);
        }
        let report = ResidueRestore {
            restored_messages,
            rendered_messages: self.messages.len() - restored_messages,
        };
        let text = if let (Some(trailing_newline), 0) = (whole_source, report.rendered_messages) {
            let mut text = lines.join("\n");
            if trailing_newline {
                text.push('\n');
            }
            text
        } else {
            let mut spliced = self.clone();
            spliced.meta.source = format.source();
            // `spliced_prefix_lens` reads the raw prefix as `raw.len()` minus one line per
            // appended message; the placeholders stand for those messages and are never written.
            lines.extend(std::iter::repeat_n(String::new(), report.rendered_messages));
            spliced.raw = lines;
            spliced.raw_is_verbatim = true;
            spliced.imported_message_count = Some(restored_messages);
            spliced.to_jsonl_spliced(format, None)?
        };
        // Soundness: what is written must read back as exactly this conversation. A segment that
        // does not (tampered, stale, or a loader that reads its own lines differently) is not used.
        let Ok(written) = Session::load_str(&text, format) else {
            return Ok(Restored::Unsound {
                first_mismatch: 0,
                restored_groups,
            });
        };
        let actual = residue_chain_keys(&written.messages);
        let first_mismatch = keyed
            .iter()
            .zip(actual.iter())
            .position(|((_, expected), (_, actual))| expected != actual)
            .unwrap_or_else(|| keyed.len().min(actual.len()));
        if first_mismatch < keyed.len() || actual.len() != keyed.len() {
            return Ok(Restored::Unsound {
                first_mismatch,
                restored_groups,
            });
        }
        Ok(Restored::Sound(text, report))
    }
}

enum Restored {
    Sound(String, ResidueRestore),
    Unsound {
        first_mismatch: usize,
        restored_groups: usize,
    },
    Nothing,
}

/// How much of a written session came from its source's segments.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResidueRestore {
    /// Messages whose source records were restored verbatim.
    pub restored_messages: usize,
    /// Messages after the restored prefix, written by the format's own writer.
    pub rendered_messages: usize,
}