1use std::collections::HashMap;
2
3use tokio::sync::broadcast;
4use tokio_util::sync::CancellationToken;
5
6use crate::error::RuntimeError;
7use crate::event::{NodeEvent, Observable};
8use crate::message::{Message, MessagePart, MessageRole};
9use crate::provider::{
10 AssistantMessage, CallTiming, DEFAULT_STREAM_BUFFER, LlmRequest, Provider, StopReason,
11 TokenUsage, estimate_tokens,
12};
13use crate::tool::BoxFut;
14use crate::value::Value;
15
16pub struct MockProvider {
17 name: String,
18 by_model: HashMap<String, Value>,
19 by_prefix: Vec<(String, String, Value)>,
20 fallback: Option<Value>,
21 chunk_delay: Option<std::time::Duration>,
22}
23
24impl MockProvider {
25 pub fn new(name: impl Into<String>) -> Self {
26 Self {
27 name: name.into(),
28 by_model: HashMap::new(),
29 by_prefix: Vec::new(),
30 fallback: None,
31 chunk_delay: None,
32 }
33 }
34
35 pub fn with_chunk_delay(mut self, d: std::time::Duration) -> Self {
36 self.chunk_delay = Some(d);
37 self
38 }
39
40 pub fn with_model(mut self, model: impl Into<String>, value: Value) -> Self {
41 self.by_model.insert(model.into(), value);
42 self
43 }
44
45 pub fn with_prefix(
46 mut self,
47 model: impl Into<String>,
48 prompt_prefix: impl Into<String>,
49 value: Value,
50 ) -> Self {
51 self.by_prefix
52 .push((model.into(), prompt_prefix.into(), value));
53 self
54 }
55
56 pub fn with_fallback(mut self, value: Value) -> Self {
57 self.fallback = Some(value);
58 self
59 }
60}
61
62impl Provider for MockProvider {
63 fn name(&self) -> &str {
64 &self.name
65 }
66
67 fn call<'a>(&'a self, req: LlmRequest) -> BoxFut<'a, Result<AssistantMessage, RuntimeError>> {
68 let turn_id = turn_id_from_req(&req);
69 Box::pin(async move {
70 self.lookup(&req)
71 .map(|v| value_to_assistant_message(&v, turn_id))
72 })
73 }
74
75 fn call_streaming(&self, req: LlmRequest) -> Observable<AssistantMessage> {
76 let turn_id = turn_id_from_req(&req);
77 let (tx, events) = broadcast::channel(DEFAULT_STREAM_BUFFER);
78 let cancel = CancellationToken::new();
79 let cancel_for_task = cancel.clone();
80 let looked_up = self.lookup(&req);
81 let chunk_delay = self.chunk_delay;
82 let output: BoxFut<'static, Result<AssistantMessage, RuntimeError>> =
83 Box::pin(async move {
84 let value = match looked_up {
85 Ok(v) => v,
86 Err(e) => {
87 let _ = tx.send(NodeEvent::LlmDone { total_tokens: 0 });
88 return Err(e);
89 }
90 };
91 let text_form = value_to_stream_text(&value);
92 let chunks = split_for_stream(&text_form);
93 let mut running = 0u64;
94 for chunk in chunks {
95 if let Some(d) = chunk_delay {
96 tokio::select! {
97 biased;
98 _ = cancel_for_task.cancelled() => {
99 let _ = tx.send(NodeEvent::LlmDone { total_tokens: running });
100 return Err(RuntimeError::Cancelled("mock stream cancelled".into()));
101 }
102 _ = tokio::time::sleep(d) => {}
103 }
104 }
105 if cancel_for_task.is_cancelled() {
106 let _ = tx.send(NodeEvent::LlmDone {
107 total_tokens: running,
108 });
109 return Err(RuntimeError::Cancelled("mock stream cancelled".into()));
110 }
111 let inc = estimate_tokens(&chunk);
112 running += inc;
113 let _ = tx.send(NodeEvent::LlmChunk {
114 text: chunk,
115 cumulative_tokens: running,
116 });
117 }
118 let _ = tx.send(NodeEvent::LlmDone {
119 total_tokens: running,
120 });
121 Ok(value_to_assistant_message(&value, turn_id))
122 });
123 Observable {
124 output,
125 events,
126 cancel,
127 }
128 }
129}
130
131impl MockProvider {
132 fn lookup(&self, req: &LlmRequest) -> Result<Value, RuntimeError> {
133 let prompt_text = req
134 .messages
135 .last()
136 .map(|m| m.text_concat())
137 .unwrap_or_default();
138 for (model, prefix, value) in &self.by_prefix {
139 if req.model == *model && prompt_text.starts_with(prefix.as_str()) {
140 return Ok(value.clone());
141 }
142 }
143 if let Some(v) = self.by_model.get(&req.model) {
144 return Ok(v.clone());
145 }
146 if let Some(v) = &self.fallback {
147 return Ok(v.clone());
148 }
149 Err(RuntimeError::ToolFailed(format!(
150 "mock provider `{}` has no entry for model={} prompt.prefix={:?}",
151 self.name,
152 req.model,
153 prompt_text.chars().take(40).collect::<String>()
154 )))
155 }
156}
157
158fn value_to_assistant_message(v: &Value, turn_id: crate::event::TurnId) -> AssistantMessage {
159 let text = value_to_stream_text(v);
160 AssistantMessage {
161 message: Message {
162 role: MessageRole::Assistant,
163 parts: vec![MessagePart::Text { text: text.clone() }],
164 turn_id,
165 },
166 stop_reason: StopReason::End,
167 token_usage: TokenUsage {
168 output: estimate_tokens(&text),
169 ..Default::default()
170 },
171 timing: CallTiming::default(),
172 model: String::new(),
173 response_id: None,
174 }
175}
176
177fn turn_id_from_req(req: &LlmRequest) -> crate::event::TurnId {
178 req.messages
179 .first()
180 .map(|m| m.turn_id.clone())
181 .unwrap_or_else(crate::event::TurnId::now)
182}
183
184fn value_to_stream_text(v: &Value) -> String {
185 match v {
186 Value::Str(s) => s.clone(),
187 other => other.to_json().to_string(),
188 }
189}
190
191fn split_for_stream(s: &str) -> Vec<String> {
192 if s.len() > 8 {
193 s.as_bytes()
194 .chunks(s.len().div_ceil(3))
195 .map(|c| String::from_utf8_lossy(c).into_owned())
196 .collect()
197 } else {
198 vec![s.to_string()]
199 }
200}
201
202#[cfg(test)]
203mod tests {
204 use super::*;
205 use crate::provider::user_text_message;
206
207 fn req(model: &str, prompt: &str) -> LlmRequest {
208 LlmRequest {
209 model: model.into(),
210 messages: vec![user_text_message(prompt)],
211 system: None,
212 input: Value::Unit,
213 schema: None,
214 cache_prompt: false,
215 tools: Vec::new(),
216 thinking_enabled: false,
217 stall_timeout_secs: 0,
218 }
219 }
220
221 #[tokio::test]
222 async fn resolves_by_model_name() {
223 let p = MockProvider::new("mock").with_model("gpt-4o-mini", Value::Str("hi".into()));
224 let out = p.call(req("gpt-4o-mini", "anything")).await.unwrap();
225 assert_eq!(out.text_concat(), "hi");
226 }
227
228 #[tokio::test]
229 async fn prefix_wins_over_model() {
230 let p = MockProvider::new("mock")
231 .with_model("m", Value::Str("model-hit".into()))
232 .with_prefix("m", "review", Value::Str("prefix-hit".into()));
233 let out = p.call(req("m", "review please")).await.unwrap();
234 assert_eq!(out.text_concat(), "prefix-hit");
235 }
236
237 #[tokio::test]
238 async fn missing_entry_errors_with_hint() {
239 let p = MockProvider::new("mock");
240 let err = p.call(req("gpt", "hello")).await.unwrap_err();
241 assert!(matches!(err, RuntimeError::ToolFailed(msg) if msg.contains("gpt")));
242 }
243
244 #[tokio::test]
245 async fn fallback_captures_unmatched() {
246 let p = MockProvider::new("mock").with_fallback(Value::Str("fb".into()));
247 let out = p.call(req("anything", "")).await.unwrap();
248 assert_eq!(out.text_concat(), "fb");
249 }
250}