1use libmcp::{
4 CompletedPendingRequest, DetailLevel, DispatchQueueOutcome, FramedMessage, HostRejection,
5 HostSessionKernel, ProbeResolution, ProbeResolutionOutcome, ReplayBudget, ReplayContract,
6 ReplayRequeueOutcome, RequestId, ToolProjection,
7};
8use serde::de::DeserializeOwned;
9use std::{
10 fs::File,
11 io::{self, BufRead, BufReader},
12 ops::Deref,
13 path::{Path, PathBuf},
14};
15
16#[derive(Debug)]
18pub struct TestCell(tempfile::TempDir);
19
20impl TestCell {
21 pub fn new(prefix: &str) -> io::Result<Self> {
23 tempfile::Builder::new().prefix(prefix).tempdir().map(Self)
24 }
25
26 pub fn new_in(prefix: &str, root: &Path) -> io::Result<Self> {
28 tempfile::Builder::new()
29 .prefix(prefix)
30 .tempdir_in(root)
31 .map(Self)
32 }
33
34 #[must_use]
36 pub fn path(&self) -> &Path {
37 self.0.path()
38 }
39
40 pub fn close(self) -> io::Result<()> {
42 self.0.close()
43 }
44
45 #[must_use]
47 pub fn persist(self) -> PathBuf {
48 self.0.keep()
49 }
50}
51
52impl AsRef<Path> for TestCell {
53 fn as_ref(&self) -> &Path {
54 self.path()
55 }
56}
57
58impl Deref for TestCell {
59 type Target = Path;
60
61 fn deref(&self) -> &Self::Target {
62 self.path()
63 }
64}
65
66#[derive(Debug, Clone)]
68pub struct ChurnHarness {
69 kernel: HostSessionKernel,
70 dispatched: Vec<FramedMessage>,
71}
72
73impl ChurnHarness {
74 #[must_use]
76 pub fn cold() -> Self {
77 Self {
78 kernel: HostSessionKernel::cold(),
79 dispatched: Vec::new(),
80 }
81 }
82
83 #[must_use]
85 pub const fn kernel(&self) -> &HostSessionKernel {
86 &self.kernel
87 }
88
89 #[must_use]
91 pub fn dispatched(&self) -> &[FramedMessage] {
92 &self.dispatched
93 }
94
95 pub fn dispatch_first(
97 &mut self,
98 frame: &FramedMessage,
99 contract: ReplayContract,
100 pending_capacity: usize,
101 ) -> Result<RequestId, HostRejection> {
102 let id = self
103 .kernel
104 .begin_request_dispatch(frame, contract, pending_capacity)?;
105 self.dispatched.push(frame.clone());
106 Ok(id)
107 }
108
109 pub fn kill_worker(&mut self, budget: ReplayBudget) -> ReplayRequeueOutcome {
111 self.kernel.requeue_pending_for_replay(budget)
112 }
113
114 pub fn dispatch_next(&mut self) -> Result<DispatchQueueOutcome, HostRejection> {
116 let outcome = self.kernel.pop_next_dispatch()?;
117 if let DispatchQueueOutcome::Replay(frame) = &outcome {
118 self.dispatched.push(frame.clone());
119 }
120 Ok(outcome)
121 }
122
123 pub fn resolve_probe(
125 &mut self,
126 request_id: &RequestId,
127 resolution: ProbeResolution,
128 max_attempts: u8,
129 ) -> Result<ProbeResolutionOutcome, HostRejection> {
130 self.kernel
131 .resolve_probe(request_id, resolution, max_attempts)
132 }
133
134 pub fn complete(
136 &mut self,
137 response: &FramedMessage,
138 ) -> Result<CompletedPendingRequest, HostRejection> {
139 self.kernel.complete_response(response)
140 }
141}
142
143pub fn read_json_lines<T>(path: &Path) -> io::Result<Vec<T>>
145where
146 T: DeserializeOwned,
147{
148 let file = File::open(path)?;
149 let reader = BufReader::new(file);
150 let mut records = Vec::new();
151 for line in reader.lines() {
152 let line = line?;
153 if line.trim().is_empty() {
154 continue;
155 }
156 let parsed = serde_json::from_str::<T>(line.as_str()).map_err(|error| {
157 io::Error::new(
158 io::ErrorKind::InvalidData,
159 format!("invalid JSONL test record: {error}"),
160 )
161 })?;
162 records.push(parsed);
163 }
164 Ok(records)
165}
166
167#[derive(Debug, Clone, PartialEq, Eq)]
169pub struct ProjectionAssertion {
170 path: String,
171 message: String,
172}
173
174impl ProjectionAssertion {
175 fn new(path: impl Into<String>, message: impl Into<String>) -> Self {
176 Self {
177 path: path.into(),
178 message: message.into(),
179 }
180 }
181}
182
183impl std::fmt::Display for ProjectionAssertion {
184 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
185 write!(formatter, "{}: {}", self.path, self.message)
186 }
187}
188
189impl std::error::Error for ProjectionAssertion {}
190
191pub fn assert_projection_doctrine<T>(projection: &T) -> Result<(), ProjectionAssertion>
193where
194 T: ToolProjection,
195{
196 let _concise = projection
197 .structured_projection(DetailLevel::Concise)
198 .map_err(|error| ProjectionAssertion::new("$concise", error.to_string()))?;
199 let _full = projection
200 .structured_projection(DetailLevel::Full)
201 .map_err(|error| ProjectionAssertion::new("$full", error.to_string()))?;
202 Ok(())
203}