1use std::collections::BTreeMap;
9use std::path::Path;
10
11use serde::{Deserialize, Serialize};
12use serde_json::Value;
13
14use sha2::{Digest, Sha256};
15
16use crate::provider::{CompletionRequest, CompletionResponse, ModelProvider, ProviderError, Usage};
17use crate::tools::{ToolError, ToolHost, ToolInvocation};
18
19pub const CASSETTE_VERSION: &str = "0.2";
21
22pub const SUPPORTED_CASSETTE_VERSIONS: &[&str] = &["0.1", "0.2"];
27
28#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
29#[serde(rename_all = "camelCase")]
30pub struct Cassette {
31 pub cassette_version: String,
32 pub agent: String,
34 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
40 pub inputs: BTreeMap<String, Value>,
41 pub interactions: Vec<Interaction>,
42 #[serde(default, skip_serializing_if = "Vec::is_empty")]
53 pub tool_calls: Vec<ToolExchange>,
54}
55
56#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
57#[serde(rename_all = "camelCase")]
58pub struct Interaction {
59 pub index: usize,
60 pub node: String,
62 pub request_digest: String,
64 pub response_type: String,
65 pub value: Value,
67 pub usage: Usage,
68 #[serde(default, skip_serializing_if = "Option::is_none")]
70 pub model: Option<String>,
71}
72
73#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
75#[serde(rename_all = "camelCase")]
76pub struct ToolExchange {
77 pub index: usize,
78 pub node: String,
80 pub tool: String,
82 pub invocation_digest: String,
84 pub result_type: String,
86 #[serde(default, skip_serializing_if = "Option::is_none")]
92 pub value: Option<Value>,
93 #[serde(default, skip_serializing_if = "Option::is_none")]
99 pub error: Option<String>,
100}
101
102impl Cassette {
103 pub fn new(agent: impl Into<String>) -> Cassette {
104 Cassette {
105 cassette_version: CASSETTE_VERSION.to_string(),
106 agent: agent.into(),
107 inputs: BTreeMap::new(),
108 interactions: Vec::new(),
109 tool_calls: Vec::new(),
110 }
111 }
112
113 pub fn to_canonical_json(&self) -> String {
118 let mut json =
119 serde_json::to_string_pretty(self).expect("a cassette is always serializable");
120 json.push('\n');
121 json
122 }
123
124 pub fn from_json(text: &str) -> Result<Cassette, String> {
125 let cassette: Cassette =
126 serde_json::from_str(text).map_err(|error| format!("invalid cassette: {error}"))?;
127 if !SUPPORTED_CASSETTE_VERSIONS.contains(&cassette.cassette_version.as_str()) {
128 return Err(format!(
129 "cassette version `{}` is not supported by this compiler (supported: {})",
130 cassette.cassette_version,
131 SUPPORTED_CASSETTE_VERSIONS.join(", ")
132 ));
133 }
134 if cassette.cassette_version == "0.1" && !cassette.tool_calls.is_empty() {
138 return Err(
139 "the cassette states version `0.1` and carries tool calls, which 0.1 has no \
140 field for; re-record it"
141 .to_string(),
142 );
143 }
144 Ok(cassette)
145 }
146
147 pub fn load(path: impl AsRef<Path>) -> Result<Cassette, String> {
148 let path = path.as_ref();
149 let text = std::fs::read_to_string(path)
150 .map_err(|error| format!("cannot read {}: {error}", path.display()))?;
151 Cassette::from_json(&text)
152 }
153
154 pub fn save(&self, path: impl AsRef<Path>) -> Result<(), String> {
155 let path = path.as_ref();
156 if let Some(parent) = path.parent() {
157 if !parent.as_os_str().is_empty() {
158 std::fs::create_dir_all(parent)
159 .map_err(|error| format!("cannot create {}: {error}", parent.display()))?;
160 }
161 }
162 std::fs::write(path, self.to_canonical_json())
163 .map_err(|error| format!("cannot write {}: {error}", path.display()))
164 }
165}
166
167pub struct ReplayProvider {
169 cassette: Cassette,
170 position: usize,
171 strict: bool,
174}
175
176impl ReplayProvider {
177 pub fn new(cassette: Cassette) -> ReplayProvider {
178 ReplayProvider {
179 cassette,
180 position: 0,
181 strict: true,
182 }
183 }
184
185 pub fn lenient(mut self) -> ReplayProvider {
186 self.strict = false;
187 self
188 }
189
190 pub fn skipping(mut self, played: usize) -> ReplayProvider {
197 self.position = played.min(self.cassette.interactions.len());
198 self
199 }
200
201 pub fn remaining(&self) -> usize {
203 self.cassette
204 .interactions
205 .len()
206 .saturating_sub(self.position)
207 }
208}
209
210impl ModelProvider for ReplayProvider {
211 fn name(&self) -> &str {
212 "replay"
213 }
214
215 fn complete(
216 &mut self,
217 request: &CompletionRequest,
218 ) -> Result<CompletionResponse, ProviderError> {
219 let Some(interaction) = self.cassette.interactions.get(self.position) else {
220 return Err(ProviderError::Cassette(format!(
221 "the cassette has {} interaction(s) but the run asked for another at node `{}`; \
222 re-record it",
223 self.cassette.interactions.len(),
224 request.node
225 )));
226 };
227 self.position += 1;
228
229 if interaction.response_type != request.response_type {
233 return Err(ProviderError::Cassette(format!(
234 "interaction {} recorded a `{}` response but node `{}` now asks for `{}`; \
235 re-record the cassette",
236 interaction.index, interaction.response_type, request.node, request.response_type
237 )));
238 }
239
240 if self.strict && interaction.request_digest != request.digest() {
241 return Err(ProviderError::Cassette(format!(
242 "interaction {} was recorded for a different request at node `{}`. \
243 The prompt or its context changed since recording — \
244 re-record the cassette and review the diff.",
245 interaction.index, request.node
246 )));
247 }
248
249 Ok(CompletionResponse {
250 value: interaction.value.clone(),
251 usage: interaction.usage,
252 model: interaction
253 .model
254 .clone()
255 .unwrap_or_else(|| "replay".to_string()),
256 })
257 }
258}
259
260pub struct RecordingProvider<P: ModelProvider> {
262 inner: P,
263 cassette: Cassette,
264}
265
266impl<P: ModelProvider> RecordingProvider<P> {
267 pub fn new(inner: P, agent: impl Into<String>) -> RecordingProvider<P> {
268 RecordingProvider {
269 inner,
270 cassette: Cassette::new(agent),
271 }
272 }
273
274 pub fn with_inputs(mut self, inputs: BTreeMap<String, Value>) -> RecordingProvider<P> {
276 self.cassette.inputs = inputs;
277 self
278 }
279
280 pub fn finish(self) -> Cassette {
281 self.cassette
282 }
283
284 fn record(&mut self, request: &CompletionRequest, response: &CompletionResponse) {
290 self.cassette.interactions.push(Interaction {
291 index: self.cassette.interactions.len(),
292 node: request.node.clone(),
293 request_digest: request.digest(),
294 response_type: request.response_type.clone(),
295 value: response.value.clone(),
296 usage: response.usage,
297 model: Some(response.model.clone()),
298 });
299 }
300}
301
302impl<P: ModelProvider> ModelProvider for RecordingProvider<P> {
303 fn name(&self) -> &str {
304 self.inner.name()
305 }
306
307 fn complete(
308 &mut self,
309 request: &CompletionRequest,
310 ) -> Result<CompletionResponse, ProviderError> {
311 let response = self.inner.complete(request)?;
312 self.record(request, &response);
313 Ok(response)
314 }
315
316 fn streams(&self) -> bool {
317 self.inner.streams()
318 }
319
320 fn complete_streaming(
321 &mut self,
322 request: &CompletionRequest,
323 on_delta: crate::provider::DeltaSink<'_>,
324 ) -> Result<CompletionResponse, ProviderError> {
325 let response = self.inner.complete_streaming(request, on_delta)?;
326 self.record(request, &response);
327 Ok(response)
328 }
329}
330
331pub fn invocation_digest(invocation: &ToolInvocation) -> String {
338 let mut hasher = Sha256::new();
339 hasher.update(invocation.agent.as_bytes());
340 hasher.update([0]);
341 hasher.update(invocation.reference.as_bytes());
342 hasher.update([0]);
343 hasher.update(invocation.result_type.as_bytes());
344 for (name, value) in &invocation.arguments {
347 hasher.update([0]);
348 hasher.update(name.as_bytes());
349 hasher.update([0]);
350 hasher.update(value.to_string().as_bytes());
351 }
352 format!("{:x}", hasher.finalize())
353}
354
355pub struct ReplayToolHost {
361 calls: Vec<ToolExchange>,
362 position: usize,
363 strict: bool,
364}
365
366impl ReplayToolHost {
367 pub fn new(calls: Vec<ToolExchange>) -> ReplayToolHost {
368 ReplayToolHost {
369 calls,
370 position: 0,
371 strict: true,
372 }
373 }
374
375 pub fn lenient(mut self) -> ReplayToolHost {
376 self.strict = false;
377 self
378 }
379
380 pub fn skipping(mut self, played: usize) -> ReplayToolHost {
383 self.position = played.min(self.calls.len());
384 self
385 }
386
387 pub fn remaining(&self) -> usize {
389 self.calls.len().saturating_sub(self.position)
390 }
391}
392
393impl ToolHost for ReplayToolHost {
394 fn name(&self) -> &str {
395 "replay"
396 }
397
398 fn provides(&self, _tool: &str) -> bool {
403 true
404 }
405
406 fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError> {
407 let Some(recorded) = self.calls.get(self.position) else {
408 return Err(ToolError::Failed(format!(
409 "the cassette records {} tool call(s) and the run asked for another: `{}`; re-record it",
410 self.calls.len(),
411 invocation.name
412 )));
413 };
414 self.position += 1;
415
416 if recorded.tool != invocation.name {
419 return Err(ToolError::Failed(format!(
420 "tool call {} recorded `{}` and the run called `{}`; re-record the cassette",
421 recorded.index, recorded.tool, invocation.name
422 )));
423 }
424 if self.strict && recorded.invocation_digest != invocation_digest(invocation) {
425 return Err(ToolError::Failed(format!(
426 "tool call {} recorded different arguments for `{}`. The call changed since recording — re-record the cassette and review the diff.",
427 recorded.index, recorded.tool
428 )));
429 }
430
431 match (&recorded.value, &recorded.error) {
432 (Some(value), _) => Ok(value.clone()),
433 (None, Some(error)) => Err(ToolError::Failed(error.clone())),
434 (None, None) => Err(ToolError::InvalidResult(format!(
435 "tool call {} for `{}` recorded neither a value nor an error",
436 recorded.index, recorded.tool
437 ))),
438 }
439 }
440}
441
442pub struct RecordingTools<H: ToolHost> {
444 inner: H,
445 calls: Vec<ToolExchange>,
446}
447
448impl<H: ToolHost> RecordingTools<H> {
449 pub fn new(inner: H) -> RecordingTools<H> {
450 RecordingTools {
451 inner,
452 calls: Vec::new(),
453 }
454 }
455
456 pub fn finish(self) -> Vec<ToolExchange> {
457 self.calls
458 }
459}
460
461impl<H: ToolHost> ToolHost for RecordingTools<H> {
462 fn name(&self) -> &str {
463 self.inner.name()
464 }
465
466 fn provides(&self, tool: &str) -> bool {
467 self.inner.provides(tool)
468 }
469
470 fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError> {
471 let result = self.inner.call(invocation);
472 let (value, error) = match &result {
473 Ok(value) => (Some(value.clone()), None),
474 Err(error) => (None, Some(error.to_string())),
478 };
479 self.calls.push(ToolExchange {
480 index: self.calls.len(),
481 node: invocation.node.clone(),
482 tool: invocation.name.clone(),
483 invocation_digest: invocation_digest(invocation),
484 result_type: invocation.result_type.clone(),
485 value,
486 error,
487 });
488 result
489 }
490}
491
492pub struct ScriptedProvider {
494 answers: Vec<Value>,
495 position: usize,
496 usage: Usage,
497}
498
499impl ScriptedProvider {
500 pub fn new(answers: Vec<Value>) -> ScriptedProvider {
501 ScriptedProvider {
502 answers,
503 position: 0,
504 usage: Usage {
505 input_tokens: 10,
506 output_tokens: 5,
507 cache_read_tokens: 0,
508 },
509 }
510 }
511
512 pub fn with_usage(mut self, usage: Usage) -> ScriptedProvider {
513 self.usage = usage;
514 self
515 }
516
517 pub fn calls(&self) -> usize {
519 self.position
520 }
521}
522
523impl ModelProvider for ScriptedProvider {
524 fn name(&self) -> &str {
525 "scripted"
526 }
527
528 fn complete(
529 &mut self,
530 request: &CompletionRequest,
531 ) -> Result<CompletionResponse, ProviderError> {
532 let Some(value) = self.answers.get(self.position).cloned() else {
533 return Err(ProviderError::Cassette(format!(
534 "the script has {} answer(s) but node `{}` asked for another",
535 self.answers.len(),
536 request.node
537 )));
538 };
539 self.position += 1;
540 Ok(CompletionResponse {
541 value,
542 usage: self.usage,
543 model: "scripted".to_string(),
544 })
545 }
546}
547
548pub fn load_directory(dir: impl AsRef<Path>) -> Result<BTreeMap<String, Cassette>, String> {
550 let dir = dir.as_ref();
551 let mut cassettes = BTreeMap::new();
552 let entries = std::fs::read_dir(dir)
553 .map_err(|error| format!("cannot read {}: {error}", dir.display()))?;
554 for entry in entries.flatten() {
555 let path = entry.path();
556 if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
557 continue;
558 }
559 let name = path
560 .file_stem()
561 .and_then(|stem| stem.to_str())
562 .unwrap_or_default()
563 .to_string();
564 cassettes.insert(name, Cassette::load(&path)?);
565 }
566 Ok(cassettes)
567}
568
569#[cfg(test)]
570mod tests {
571 use super::*;
572 use crate::schema::ResponseShape;
573 use serde_json::json;
574
575 fn request(node: &str, prompt: &str) -> CompletionRequest {
576 CompletionRequest {
577 node: node.into(),
578 model: crate::provider::ModelSelection::Default,
579 system: None,
580 prompt: prompt.into(),
581 context: Vec::new(),
582 response_type: "markdown".into(),
583 shape: ResponseShape::Prose,
584 max_tokens: 1024,
585 }
586 }
587
588 fn recorded() -> Cassette {
589 let mut provider = RecordingProvider::new(
590 ScriptedProvider::new(vec![json!("first"), json!("second")]),
591 "test.Agent",
592 );
593 provider.complete(&request("n0", "one")).unwrap();
594 provider.complete(&request("n1", "two")).unwrap();
595 provider.finish()
596 }
597
598 #[test]
599 fn a_cassette_round_trips_through_canonical_json() {
600 let cassette = recorded();
601 let parsed = Cassette::from_json(&cassette.to_canonical_json()).unwrap();
602 assert_eq!(parsed, cassette);
603 }
604
605 #[test]
606 fn canonical_json_ends_with_a_newline() {
607 assert!(recorded().to_canonical_json().ends_with("}\n"));
608 }
609
610 #[test]
611 fn replay_serves_recorded_answers_in_order() {
612 let mut provider = ReplayProvider::new(recorded());
613 assert_eq!(
614 provider.complete(&request("n0", "one")).unwrap().value,
615 json!("first")
616 );
617 assert_eq!(
618 provider.complete(&request("n1", "two")).unwrap().value,
619 json!("second")
620 );
621 assert_eq!(provider.remaining(), 0);
622 }
623
624 #[test]
625 fn replay_rejects_a_changed_prompt() {
626 let mut provider = ReplayProvider::new(recorded());
627 let error = provider
628 .complete(&request("n0", "a different prompt"))
629 .unwrap_err();
630 let message = error.to_string();
631 assert!(
632 message.contains("recorded for a different request"),
633 "{message}"
634 );
635 assert!(message.contains("re-record"), "{message}");
636 }
637
638 #[test]
639 fn replay_rejects_a_changed_response_type() {
640 let mut provider = ReplayProvider::new(recorded());
641 let mut changed = request("n0", "one");
642 changed.response_type = "string".into();
643 let error = provider.complete(&changed).unwrap_err();
644 assert!(
645 error.to_string().contains("recorded a `markdown` response"),
646 "{error}"
647 );
648 }
649
650 #[test]
651 fn replay_reports_an_exhausted_cassette() {
652 let mut provider = ReplayProvider::new(recorded());
653 provider.complete(&request("n0", "one")).unwrap();
654 provider.complete(&request("n1", "two")).unwrap();
655 let error = provider.complete(&request("n2", "three")).unwrap_err();
656 assert!(error.to_string().contains("asked for another"), "{error}");
657 }
658
659 fn tool_call(name: &str, path: &str) -> ToolInvocation {
662 ToolInvocation {
663 node: "n0".into(),
664 agent: "test.Agent".into(),
665 reference: format!("mcp:{name}"),
666 name: name.into(),
667 transport: "mcp".into(),
668 arguments: [("path".to_string(), json!(path))].into(),
669 effects: vec!["filesystem_read".into()],
670 result_type: "text".into(),
671 }
672 }
673
674 struct FixedTool(Result<Value, &'static str>);
675
676 impl ToolHost for FixedTool {
677 fn name(&self) -> &str {
678 "fixed"
679 }
680 fn provides(&self, _tool: &str) -> bool {
681 true
682 }
683 fn call(&mut self, _invocation: &ToolInvocation) -> Result<Value, ToolError> {
684 self.0
685 .clone()
686 .map_err(|error| ToolError::Failed(error.into()))
687 }
688 }
689
690 #[test]
691 fn a_recorded_tool_call_replays_without_reaching_anything() {
692 let mut recorder = RecordingTools::new(FixedTool(Ok(json!("# Sample"))));
693 recorder
694 .call(&tool_call("fs.read_file", "README.md"))
695 .unwrap();
696 let recorded = recorder.finish();
697 assert_eq!(recorded.len(), 1);
698 assert_eq!(recorded[0].tool, "fs.read_file");
699
700 let mut replay = ReplayToolHost::new(recorded);
701 assert_eq!(
702 replay
703 .call(&tool_call("fs.read_file", "README.md"))
704 .unwrap(),
705 json!("# Sample")
706 );
707 assert_eq!(replay.remaining(), 0);
708 }
709
710 #[test]
711 fn a_recorded_failure_replays_as_a_failure() {
712 let mut recorder = RecordingTools::new(FixedTool(Err("no such file")));
716 assert!(recorder
717 .call(&tool_call("fs.read_file", "gone.md"))
718 .is_err());
719 let recorded = recorder.finish();
720 assert_eq!(recorded[0].value, None);
721 assert!(recorded[0]
722 .error
723 .as_deref()
724 .unwrap()
725 .contains("no such file"));
726
727 let mut replay = ReplayToolHost::new(recorded);
728 let error = replay
729 .call(&tool_call("fs.read_file", "gone.md"))
730 .unwrap_err();
731 assert!(error.to_string().contains("no such file"), "{error}");
732 }
733
734 #[test]
735 fn replay_refuses_a_call_whose_arguments_changed() {
736 let mut recorder = RecordingTools::new(FixedTool(Ok(json!("# Sample"))));
737 recorder
738 .call(&tool_call("fs.read_file", "README.md"))
739 .unwrap();
740 let mut replay = ReplayToolHost::new(recorder.finish());
741
742 let error = replay
743 .call(&tool_call("fs.read_file", "notes.md"))
744 .unwrap_err();
745 assert!(error.to_string().contains("re-record"), "{error}");
746 }
747
748 #[test]
749 fn replay_refuses_a_different_tool_by_name() {
750 let mut recorder = RecordingTools::new(FixedTool(Ok(json!("# Sample"))));
751 recorder
752 .call(&tool_call("fs.read_file", "README.md"))
753 .unwrap();
754 let mut replay = ReplayToolHost::new(recorder.finish());
755
756 let error = replay
757 .call(&tool_call("fs.list_dir", "README.md"))
758 .unwrap_err();
759 let message = error.to_string();
760 assert!(message.contains("fs.read_file"), "{message}");
761 assert!(message.contains("fs.list_dir"), "{message}");
762 }
763
764 #[test]
765 fn replay_reports_a_call_beyond_the_recording() {
766 let mut replay = ReplayToolHost::new(Vec::new());
767 let error = replay
768 .call(&tool_call("fs.read_file", "README.md"))
769 .unwrap_err();
770 assert!(error.to_string().contains("asked for another"), "{error}");
771 }
772
773 #[test]
774 fn the_invocation_digest_ignores_effects_and_notices_the_agent() {
775 let base = tool_call("fs.read_file", "README.md");
776
777 let mut other_effects = tool_call("fs.read_file", "README.md");
778 other_effects.effects = vec!["filesystem_write".into()];
779 assert_eq!(
780 invocation_digest(&base),
781 invocation_digest(&other_effects),
782 "effects say what a call may do, not what it answers"
783 );
784
785 let mut other_agent = tool_call("fs.read_file", "README.md");
786 other_agent.agent = "test.Other".into();
787 assert_ne!(
788 invocation_digest(&base),
789 invocation_digest(&other_agent),
790 "two agents hold different policies, so the same call from another is another call"
791 );
792 }
793
794 #[test]
795 fn a_zero_one_cassette_still_replays_and_a_lying_one_does_not() {
796 let mut cassette = recorded();
797 cassette.cassette_version = "0.1".into();
798 let parsed = Cassette::from_json(&cassette.to_canonical_json()).unwrap();
799 assert!(parsed.tool_calls.is_empty());
800
801 cassette.tool_calls.push(ToolExchange {
802 index: 0,
803 node: "n0".into(),
804 tool: "fs.read_file".into(),
805 invocation_digest: "x".into(),
806 result_type: "text".into(),
807 value: Some(json!("hi")),
808 error: None,
809 });
810 let error = Cassette::from_json(&cassette.to_canonical_json()).unwrap_err();
811 assert!(error.contains("re-record"), "{error}");
812 }
813
814 #[test]
815 fn a_future_cassette_version_is_rejected() {
816 let mut cassette = recorded();
817 cassette.cassette_version = "9.0".into();
818 let error = Cassette::from_json(&cassette.to_canonical_json()).unwrap_err();
819 assert!(error.contains("not supported"), "{error}");
820 }
821
822 #[test]
823 fn lenient_replay_tolerates_a_changed_prompt() {
824 let mut provider = ReplayProvider::new(recorded()).lenient();
825 assert!(provider.complete(&request("n0", "changed")).is_ok());
826 }
827}