atomr_agents_agent/
replay.rs

1//! Side-effect-free replay (FR-1).
2//!
3//! [`ReplayProvider`] implements [`InferenceClient`] by serving the
4//! *recorded* completion for each step instead of calling a live
5//! provider. Construct an [`Agent`](crate::Agent) with a `ReplayProvider`
6//! as its inference client and the model side of a run reproduces
7//! byte-identically with **zero provider calls**. A missing record is a
8//! loud, typed [`ReplayError`] — never a silent fall-back to live
9//! inference.
10//!
11//! The companion [`record_turn`] converts a live [`TurnResult`] into an
12//! [`InferenceRecord`] so a recording run can persist provenance through
13//! a [`RecordingCheckpointer`](atomr_agents_state::RecordingCheckpointer).
14
15use async_trait::async_trait;
16use atomr_agents_core::{AgentError, Result, RunId, Value, WorkflowId};
17use atomr_agents_state::{InferenceRecord, StepRecord, StepRecordStore, UsageRecord};
18use atomr_agents_tool::{ParsedToolCall, Provider};
19use atomr_infer_core::batch::ExecuteBatch;
20use atomr_infer_core::tokens::{FinishReason, TokenUsage};
21use parking_lot::Mutex;
22use thiserror::Error;
23
24use crate::inference::{InferenceClient, TurnResult};
25
26/// Errors raised on the replay path. Replay fails loudly rather than
27/// re-inferring.
28#[derive(Debug, Error)]
29pub enum ReplayError {
30    #[error("replay: no record for step {step} of run {run} (have {available} steps)")]
31    MissingStep {
32        run: String,
33        step: usize,
34        available: usize,
35    },
36    #[error("replay: step {step} has no recorded inference")]
37    NoInference { step: usize },
38    #[error("replay: could not decode recorded value: {0}")]
39    Decode(String),
40}
41
42impl From<ReplayError> for AgentError {
43    fn from(e: ReplayError) -> Self {
44        AgentError::Inference(e.to_string())
45    }
46}
47
48/// An [`InferenceClient`] that returns recorded completions in order.
49pub struct ReplayProvider {
50    provider: Provider,
51    run: String,
52    steps: Vec<StepRecord>,
53    cursor: Mutex<usize>,
54}
55
56impl ReplayProvider {
57    /// Build a replay provider from recorded steps (ordered ascending by
58    /// `super_step`). The provider discriminant is taken from the first
59    /// recorded inference (defaulting to OpenAI if none is present —
60    /// irrelevant for replay since no parsing of live deltas occurs).
61    pub fn from_steps(run: impl Into<String>, steps: Vec<StepRecord>) -> Self {
62        let provider = steps
63            .iter()
64            .find_map(|s| s.inference.as_ref())
65            .and_then(|inf| provider_from_str(&inf.provider))
66            .unwrap_or(Provider::OpenAi);
67        Self {
68            provider,
69            run: run.into(),
70            steps,
71            cursor: Mutex::new(0),
72        }
73    }
74
75    /// Load recorded steps for a run from a [`StepRecordStore`].
76    pub async fn load(store: &dyn StepRecordStore, workflow: &WorkflowId, run: &RunId) -> Result<Self> {
77        let steps = store.list_records(workflow, run).await?;
78        Ok(Self::from_steps(run.as_str(), steps))
79    }
80
81    /// Number of recorded steps remaining to replay.
82    pub fn remaining(&self) -> usize {
83        self.steps.len().saturating_sub(*self.cursor.lock())
84    }
85
86    fn next_record(&self) -> std::result::Result<TurnResult, ReplayError> {
87        let mut cur = self.cursor.lock();
88        let idx = *cur;
89        let step = self.steps.get(idx).ok_or(ReplayError::MissingStep {
90            run: self.run.clone(),
91            step: idx,
92            available: self.steps.len(),
93        })?;
94        let inf = step
95            .inference
96            .as_ref()
97            .ok_or(ReplayError::NoInference { step: idx })?;
98        let turn = turn_from_record(inf)?;
99        *cur += 1;
100        Ok(turn)
101    }
102}
103
104#[async_trait]
105impl InferenceClient for ReplayProvider {
106    fn provider(&self) -> Provider {
107        self.provider
108    }
109
110    async fn run(&self, _batch: ExecuteBatch) -> Result<TurnResult> {
111        // The batch is intentionally ignored: replay serves the recorded
112        // completion and never calls a provider.
113        Ok(self.next_record()?)
114    }
115}
116
117/// Reconstruct a [`TurnResult`] from a recorded [`InferenceRecord`].
118pub fn turn_from_record(inf: &InferenceRecord) -> std::result::Result<TurnResult, ReplayError> {
119    let usage = TokenUsage {
120        input_tokens: inf.usage.prompt_tokens,
121        output_tokens: inf.usage.completion_tokens,
122        reasoning_tokens: inf.usage.reasoning_tokens,
123        cached_tokens: inf.usage.cached_tokens,
124    };
125    let finish_reason = match &inf.finish_reason {
126        Some(s) => Some(
127            serde_json::from_value::<FinishReason>(Value::String(s.clone()))
128                .map_err(|e| ReplayError::Decode(format!("finish_reason '{s}': {e}")))?,
129        ),
130        None => None,
131    };
132    let tool_calls = inf
133        .tool_calls
134        .iter()
135        .map(|v| {
136            serde_json::from_value::<ParsedToolCall>(v.clone())
137                .map_err(|e| ReplayError::Decode(format!("tool_call: {e}")))
138        })
139        .collect::<std::result::Result<Vec<_>, _>>()?;
140    Ok(TurnResult {
141        text: inf.raw_completion.clone(),
142        usage,
143        finish_reason,
144        tool_calls,
145    })
146}
147
148/// Capture a live [`TurnResult`] as an [`InferenceRecord`] for later
149/// replay. `params_hash` / `prompt_hash` are supplied by the caller (they
150/// derive from the request, which the pipeline holds).
151pub fn record_turn(
152    provider: Provider,
153    model_id: impl Into<String>,
154    model_version: impl Into<String>,
155    params_hash: impl Into<String>,
156    prompt_hash: impl Into<String>,
157    turn: &TurnResult,
158) -> InferenceRecord {
159    InferenceRecord {
160        provider: provider_to_str(provider).to_string(),
161        model_id: model_id.into(),
162        model_version: model_version.into(),
163        params_hash: params_hash.into(),
164        prompt_hash: prompt_hash.into(),
165        raw_completion: turn.text.clone(),
166        finish_reason: turn
167            .finish_reason
168            .and_then(|r| serde_json::to_value(r).ok())
169            .and_then(|v| v.as_str().map(|s| s.to_string())),
170        usage: UsageRecord {
171            prompt_tokens: turn.usage.input_tokens,
172            completion_tokens: turn.usage.output_tokens,
173            reasoning_tokens: turn.usage.reasoning_tokens,
174            cached_tokens: turn.usage.cached_tokens,
175        },
176        tool_calls: turn
177            .tool_calls
178            .iter()
179            .filter_map(|tc| serde_json::to_value(tc).ok())
180            .collect(),
181    }
182}
183
184fn provider_to_str(p: Provider) -> &'static str {
185    match p {
186        Provider::OpenAi => "open_ai",
187        Provider::Anthropic => "anthropic",
188    }
189}
190
191fn provider_from_str(s: &str) -> Option<Provider> {
192    serde_json::from_value::<Provider>(Value::String(s.to_string())).ok()
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use atomr_agents_state::CheckpointKey;
199
200    fn batch() -> ExecuteBatch {
201        ExecuteBatch {
202            request_id: "req".into(),
203            model: "m".into(),
204            messages: vec![],
205            sampling: Default::default(),
206            stream: false,
207            estimated_tokens: 0,
208        }
209    }
210
211    fn rec_step(run: &str, step: u64, text: &str, with_tool: bool) -> StepRecord {
212        let tool_calls = if with_tool {
213            vec![serde_json::to_value(ParsedToolCall {
214                id: "c1".into(),
215                name: "place_order".into(),
216                arguments_raw: "{\"qty\":1}".into(),
217            })
218            .unwrap()]
219        } else {
220            vec![]
221        };
222        StepRecord::new(CheckpointKey {
223            workflow_id: WorkflowId::from("wf"),
224            run_id: RunId::from(run),
225            super_step: step,
226        })
227        .with_inference(InferenceRecord {
228            provider: "anthropic".into(),
229            model_id: "claude".into(),
230            model_version: "1".into(),
231            params_hash: "p".into(),
232            prompt_hash: "h".into(),
233            raw_completion: text.into(),
234            finish_reason: Some(if with_tool {
235                "tool_calls".into()
236            } else {
237                "stop".into()
238            }),
239            usage: UsageRecord {
240                prompt_tokens: 3,
241                completion_tokens: 2,
242                ..Default::default()
243            },
244            tool_calls,
245        })
246    }
247
248    #[tokio::test]
249    async fn replays_recorded_completions_in_order() {
250        let steps = vec![rec_step("r", 0, "first", true), rec_step("r", 1, "second", false)];
251        let rp = ReplayProvider::from_steps("r", steps);
252        assert_eq!(rp.provider(), Provider::Anthropic);
253        assert_eq!(rp.remaining(), 2);
254
255        let t0 = rp.run(batch()).await.unwrap();
256        assert_eq!(t0.text, "first");
257        assert_eq!(t0.tool_calls.len(), 1);
258        assert_eq!(t0.tool_calls[0].name, "place_order");
259        assert_eq!(t0.finish_reason, Some(FinishReason::ToolCalls));
260
261        let t1 = rp.run(batch()).await.unwrap();
262        assert_eq!(t1.text, "second");
263        assert!(t1.tool_calls.is_empty());
264        assert_eq!(rp.remaining(), 0);
265    }
266
267    #[tokio::test]
268    async fn missing_record_is_loud_error() {
269        let rp = ReplayProvider::from_steps("r", vec![rec_step("r", 0, "only", false)]);
270        let _ = rp.run(batch()).await.unwrap();
271        // Second call: no more records -> typed error, never live inference.
272        let err = rp.run(batch()).await.unwrap_err();
273        assert!(matches!(err, AgentError::Inference(_)));
274        assert!(err.to_string().contains("replay"));
275    }
276
277    #[test]
278    fn record_turn_then_replay_is_identical() {
279        let turn = TurnResult {
280            text: "hello".into(),
281            usage: TokenUsage {
282                input_tokens: 5,
283                output_tokens: 7,
284                reasoning_tokens: 1,
285                cached_tokens: 2,
286            },
287            finish_reason: Some(FinishReason::Stop),
288            tool_calls: vec![ParsedToolCall {
289                id: "x".into(),
290                name: "f".into(),
291                arguments_raw: "{}".into(),
292            }],
293        };
294        let rec = record_turn(Provider::Anthropic, "claude", "1", "p", "h", &turn);
295        let back = turn_from_record(&rec).unwrap();
296        assert_eq!(back.text, turn.text);
297        assert_eq!(back.usage.input_tokens, 5);
298        assert_eq!(back.usage.cached_tokens, 2);
299        assert_eq!(back.finish_reason, Some(FinishReason::Stop));
300        assert_eq!(back.tool_calls, turn.tool_calls);
301    }
302}
atomr_agents_agent/replay.rs

atomr_agents_agent/
replay.rs
