af_agent_runtime/

history.rs

use af_agent_session::{ContentBlock, DeliveryMode, Event, SessionEvent, SessionProjection};
use af_context::RunId;
use af_llm::{ChatMessage, InputImage};

use crate::RuntimeError;

/// Validated execution facts and the model-visible tail of one durable Run.
#[derive(Debug, Clone)]
pub struct RunHistory {
    pub(crate) run_id: RunId,
    pub(crate) projection: SessionProjection,
    pub(crate) events: Vec<SessionEvent>,
    pub(crate) transcript: Vec<ChatMessage>,
}

/// Incremental two-pass replay. First fold every fact; then read model events
/// only after the latest successful summary. No persisted checkpoint is created.
pub struct RunHistoryReplay {
    run_id: RunId,
    projection: SessionProjection,
    events: Vec<SessionEvent>,
    capturing_run: bool,
    latest_images: Vec<InputImage>,
    pending_summary: Option<(u64, String, Vec<InputImage>)>,
    summary: Option<(u64, String, Vec<InputImage>)>,
    model_seq: Option<u64>,
    transcript: Vec<ChatMessage>,
}

impl RunHistoryReplay {
    /// Start reducing a Run from the first Session event.
    pub fn new(run_id: RunId) -> Self {
        Self {
            run_id,
            projection: Default::default(),
            events: Vec::new(),
            capturing_run: false,
            latest_images: Vec::new(),
            pending_summary: None,
            summary: None,
            model_seq: None,
            transcript: Vec::new(),
        }
    }

    /// Validate and fold one bounded page of complete durable facts in order.
    pub fn facts(&mut self, events: &[SessionEvent]) -> Result<(), RuntimeError> {
        if self.model_seq.is_some() {
            return Err(RuntimeError::Invariant(
                "facts after model replay began".into(),
            ));
        }
        for event in events {
            let open = self.projection.open_compaction.clone();
            let active = self.projection.active_run_id.clone();
            self.projection
                .apply_facts(event)
                .map_err(|error| RuntimeError::Invariant(error.to_string()))?;
            match &event.event {
                Event::InputQueued {
                    run_id,
                    mode: DeliveryMode::Followup,
                    ..
                }
                | Event::RunStarted { run_id, .. }
                    if run_id == &self.run_id =>
                {
                    self.capturing_run = true
                }
                Event::UserMessage { content, .. } => {
                    self.latest_images.clear();
                    for block in content {
                        if let ContentBlock::Resource {
                            resource_id,
                            media_type,
                        } = block
                        {
                            if media_type.starts_with("image/") {
                                let image = InputImage {
                                    asset_id: resource_id.parse().map_err(|_| {
                                        RuntimeError::InvalidInput(
                                            "invalid stored image identity".into(),
                                        )
                                    })?,
                                    media_type: media_type.clone(),
                                };
                                image.validate().map_err(|error| {
                                    RuntimeError::InvalidInput(error.to_string())
                                })?;
                                self.latest_images.push(image);
                            }
                        }
                    }
                }
                Event::SummaryReplaced {
                    run_id,
                    through_seq,
                    summary,
                    ..
                } if active.as_ref() == Some(run_id)
                    && open.as_ref().is_some_and(|(_, seq)| seq == through_seq)
                    && *through_seq < event.seq =>
                {
                    self.pending_summary =
                        Some((event.seq, summary.clone(), self.latest_images.clone()));
                }
                Event::CompactionFinished { status, error, .. } => {
                    if status == "completed" && error.is_none() {
                        if let Some(summary) = self.pending_summary.take() {
                            self.summary = Some(summary);
                        }
                    } else {
                        self.pending_summary = None;
                    }
                }
                Event::CompactionStarted { .. } => self.pending_summary = None,
                _ => {}
            }
            if self.capturing_run {
                self.events.push(event.clone());
            }
        }
        Ok(())
    }

    /// Freeze the first pass and return the exclusive cursor for the model tail.
    pub fn begin_model_tail(&mut self) -> u64 {
        if let Some(seq) = self.model_seq {
            return seq;
        }
        let seq = if let Some((seq, summary, images)) = &self.summary {
            if !images.is_empty() {
                let mut message = ChatMessage::user("");
                message.images = images.clone();
                self.transcript.push(message);
            }
            crate::replay::replace_with_summary(&mut self.transcript, summary);
            *seq
        } else {
            0
        };
        self.model_seq = Some(seq);
        seq
    }

    /// Fold subsequent model events through the exact first-pass upper boundary.
    pub fn model_tail(&mut self, events: &[SessionEvent]) -> Result<(), RuntimeError> {
        let mut seq = self
            .model_seq
            .ok_or_else(|| RuntimeError::Invariant("model tail was not initialized".into()))?;
        for event in events {
            if event.seq != seq + 1
                || event.seq > self.projection.last_seq
                || self.projection.session_id.as_ref() != Some(&event.session_id)
            {
                return Err(RuntimeError::Invariant(
                    "model tail boundary changed".into(),
                ));
            }
            // Only the successful summary selected during fact validation is used.
            if !matches!(event.event, Event::SummaryReplaced { .. }) {
                crate::replay::apply_transcript_event(&mut self.transcript, event)?;
            }
            seq = event.seq;
        }
        self.model_seq = Some(seq);
        Ok(())
    }

    /// Finish only when the model pass reaches the validated facts' upper boundary.
    pub fn finish(self) -> Result<RunHistory, RuntimeError> {
        if self.model_seq != Some(self.projection.last_seq) {
            return Err(RuntimeError::Invariant("incomplete model tail".into()));
        }
        Ok(RunHistory {
            run_id: self.run_id,
            projection: self.projection,
            events: self.events,
            transcript: self.transcript,
        })
    }
}

impl RunHistory {
    /// Last fact sequence covered by both replay passes.
    pub fn last_seq(&self) -> u64 {
        self.projection.last_seq
    }

    /// Convenience for in-memory callers; stores should feed bounded pages to RunHistoryReplay.
    pub fn from_events(run_id: RunId, events: &[SessionEvent]) -> Result<Self, RuntimeError> {
        let mut replay = RunHistoryReplay::new(run_id);
        replay.facts(events)?;
        let after = replay.begin_model_tail();
        replay.model_tail(&events[events.partition_point(|event| event.seq <= after)..])?;
        replay.finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use af_agent_session::{text, RunStatus};

    fn history(completed: bool) -> Vec<SessionEvent> {
        let run_id: RunId = "old-run".parse().unwrap();
        let events = vec![
            Event::SessionCreated {
                profile_revision_id: "profile".parse().unwrap(),
            },
            Event::InputQueued {
                input_id: "old-input".parse().unwrap(),
                run_id: run_id.clone(),
                mode: DeliveryMode::Followup,
                content: text("old"),
                explicit_skill: None,
            },
            Event::InputClaimed {
                input_id: "old-input".parse().unwrap(),
                run_id: run_id.clone(),
            },
            Event::RunStarted {
                input_id: "old-input".parse().unwrap(),
                run_id: run_id.clone(),
            },
            Event::TurnStarted {
                run_id: run_id.clone(),
                turn: 1,
            },
            Event::StepStarted {
                run_id: run_id.clone(),
                step: 1,
            },
            Event::UserMessage {
                run_id: run_id.clone(),
                content: vec![
                    ContentBlock::Text {
                        text: "old large transcript".repeat(1000),
                    },
                    ContentBlock::Resource {
                        resource_id: "image".into(),
                        media_type: "image/png".into(),
                    },
                ],
            },
            Event::UsageRecorded {
                metering: None,
                run_id: run_id.clone(),
                operation_id: "attempt".into(),
                prompt_tokens: 5,
                completion_tokens: 3,
                cost_units: 0,
            },
            Event::CompactionStarted {
                run_id: run_id.clone(),
                compaction_id: "compact".into(),
                source_through_seq: 8,
            },
            Event::SummaryReplaced {
                run_id: run_id.clone(),
                through_seq: 8,
                summary: "confirmed summary".into(),
                compactor: "test".into(),
                model: "test".into(),
            },
            Event::CompactionFinished {
                run_id: run_id.clone(),
                compaction_id: "compact".into(),
                status: if completed { "completed" } else { "failed" }.into(),
                error: (!completed).then(|| "failed".into()),
            },
            Event::StepFinished {
                run_id: run_id.clone(),
                step: 1,
            },
            Event::TurnFinished {
                run_id: run_id.clone(),
                turn: 1,
            },
            Event::RunFinished {
                run_id,
                status: RunStatus::Completed,
                error_code: None,
            },
            Event::InputQueued {
                input_id: "input".parse().unwrap(),
                run_id: "new-run".parse().unwrap(),
                mode: DeliveryMode::Followup,
                content: text("new"),
                explicit_skill: None,
            },
        ];
        events
            .into_iter()
            .enumerate()
            .map(|(index, event)| SessionEvent {
                session_id: "session".parse().unwrap(),
                seq: index as u64 + 1,
                occurred_at: chrono::Utc::now(),
                event,
            })
            .collect()
    }

    #[test]
    fn successful_summary_starts_model_tail_without_losing_images_or_usage() {
        let events = history(true);
        let mut replay = RunHistoryReplay::new("new-run".parse().unwrap());
        for page in events.chunks(3) {
            replay.facts(page).unwrap();
        }
        assert_eq!(replay.begin_model_tail(), 10);
        for page in events[10..].chunks(2) {
            replay.model_tail(page).unwrap();
        }
        let resumed = replay.finish().unwrap();
        assert_eq!(resumed.projection.usage_for("old-run"), (5, 3));
        assert_eq!(resumed.events.len(), 1);
        assert_eq!(resumed.transcript.len(), 2);
        assert_eq!(
            resumed.transcript[0].content.as_deref(),
            Some("Conversation summary:\nconfirmed summary")
        );
        assert_eq!(resumed.transcript[1].images[0].asset_id, "image");
        assert!(resumed.projection.messages.is_empty());
    }

    #[test]
    fn failed_or_unfinished_summaries_never_hide_original_messages() {
        for events in [history(false), history(true)[..10].to_vec()] {
            let mut replay = RunHistoryReplay::new("new-run".parse().unwrap());
            replay.facts(&events).unwrap();
            assert_eq!(replay.begin_model_tail(), 0);
            replay.model_tail(&events).unwrap();
            let resumed = replay.finish().unwrap();
            assert!(resumed.transcript[0]
                .content
                .as_deref()
                .unwrap()
                .contains("old large transcript"));
            assert_eq!(resumed.projection.usage_for("old-run"), (5, 3));
        }
    }

    #[test]
    fn a_model_tail_cannot_skip_events_or_cross_the_validated_boundary() {
        let events = history(true);
        let mut replay = RunHistoryReplay::new("new-run".parse().unwrap());
        replay.facts(&events).unwrap();
        assert!(replay.model_tail(&events).is_err());
        replay.begin_model_tail();
        assert!(replay.model_tail(&events[11..]).is_err());
        assert!(replay.facts(&[]).is_err());
        assert!(replay.finish().is_err());
    }
}