1use std::cell::RefCell;
49use std::time::Duration;
50
51use crate::llm::api::{DeltaSender, LlmRequestPayload, LlmResult, ProviderTelemetry};
52use crate::llm::provider::{LlmProvider, LlmProviderChat};
53use crate::value::{ErrorCategory, VmError};
54
55#[derive(Clone, Debug)]
57pub enum FakeLlmTurn {
58 Stream(Vec<FakeLlmEvent>),
61 Error(FakeLlmError),
64 Stalled(Duration),
68}
69
70impl FakeLlmTurn {
71 pub fn stream(events: impl IntoIterator<Item = FakeLlmEvent>) -> Self {
73 Self::Stream(events.into_iter().collect())
74 }
75
76 pub fn error(category: ErrorCategory, message: impl Into<String>) -> Self {
78 Self::Error(FakeLlmError {
79 category,
80 message: message.into(),
81 retry_after_ms: None,
82 })
83 }
84}
85
86#[derive(Clone, Debug)]
88pub enum FakeLlmEvent {
89 Token(String),
92 ToolCallDelta {
95 id: String,
96 name: String,
97 arguments: serde_json::Value,
98 },
99 Stall(Duration),
102 Error(FakeLlmError),
105 Done(FakeStopReason),
108}
109
110#[derive(Clone, Debug, PartialEq, Eq)]
112pub enum FakeStopReason {
113 EndTurn,
114 ToolUse,
115 MaxTokens,
116 StopSequence,
117 Custom(String),
118}
119
120impl FakeStopReason {
121 fn as_str(&self) -> &str {
122 match self {
123 Self::EndTurn => "end_turn",
124 Self::ToolUse => "tool_use",
125 Self::MaxTokens => "max_tokens",
126 Self::StopSequence => "stop_sequence",
127 Self::Custom(value) => value.as_str(),
128 }
129 }
130}
131
132#[derive(Clone, Debug)]
134pub struct FakeLlmError {
135 pub category: ErrorCategory,
136 pub message: String,
137 pub retry_after_ms: Option<u64>,
141}
142
143impl FakeLlmError {
144 pub fn new(category: ErrorCategory, message: impl Into<String>) -> Self {
145 Self {
146 category,
147 message: message.into(),
148 retry_after_ms: None,
149 }
150 }
151
152 pub fn with_retry_after_ms(mut self, ms: u64) -> Self {
153 self.retry_after_ms = Some(ms);
154 self
155 }
156}
157
158#[derive(Clone, Debug, Default)]
165pub struct FakeLlmScript {
166 pub turns: Vec<FakeLlmTurn>,
167}
168
169impl FakeLlmScript {
170 pub fn new() -> Self {
171 Self::default()
172 }
173
174 pub fn streaming(events: impl IntoIterator<Item = FakeLlmEvent>) -> Self {
176 Self {
177 turns: vec![FakeLlmTurn::stream(events)],
178 }
179 }
180
181 pub fn erroring(category: ErrorCategory, message: impl Into<String>) -> Self {
183 Self {
184 turns: vec![FakeLlmTurn::error(category, message)],
185 }
186 }
187
188 pub fn push(mut self, turn: FakeLlmTurn) -> Self {
190 self.turns.push(turn);
191 self
192 }
193}
194
195#[derive(Clone, Debug)]
198pub struct FakeLlmCall {
199 pub provider: String,
200 pub model: String,
201 pub system: Option<String>,
202 pub messages: Vec<serde_json::Value>,
203 pub native_tools: Option<Vec<serde_json::Value>>,
204 pub stream: bool,
205}
206
207impl FakeLlmCall {
208 fn from_request(request: &LlmRequestPayload) -> Self {
209 Self {
210 provider: request.provider.clone(),
211 model: request.model.clone(),
212 system: request.system.clone(),
213 messages: request.messages.clone(),
214 native_tools: request.native_tools.clone(),
215 stream: request.stream,
216 }
217 }
218}
219
220thread_local! {
221 static FAKE_LLM_TURNS: RefCell<Vec<FakeLlmTurn>> = const { RefCell::new(Vec::new()) };
222 static FAKE_LLM_CALLS: RefCell<Vec<FakeLlmCall>> = const { RefCell::new(Vec::new()) };
223}
224
225#[must_use = "FakeLlmGuard asserts on drop; bind it to a `_guard` local"]
233pub fn install_fake_llm_script(script: FakeLlmScript) -> FakeLlmGuard {
234 FAKE_LLM_TURNS.with(|turns| {
235 let mut turns = turns.borrow_mut();
236 assert!(
237 turns.is_empty(),
238 "FakeLlmProvider: a script is already installed; drop the previous guard before installing a new one"
239 );
240 *turns = script.turns;
241 });
242 FAKE_LLM_CALLS.with(|calls| calls.borrow_mut().clear());
243 FakeLlmGuard { _priv: () }
244}
245
246pub fn fake_llm_captured_calls() -> Vec<FakeLlmCall> {
248 FAKE_LLM_CALLS.with(|calls| calls.borrow().clone())
249}
250
251#[must_use]
257pub struct FakeLlmGuard {
258 _priv: (),
259}
260
261impl Drop for FakeLlmGuard {
262 fn drop(&mut self) {
263 let remaining = FAKE_LLM_TURNS.with(|turns| std::mem::take(&mut *turns.borrow_mut()));
264 FAKE_LLM_CALLS.with(|calls| calls.borrow_mut().clear());
265 if std::thread::panicking() {
268 return;
269 }
270 assert!(
271 remaining.is_empty(),
272 "FakeLlmProvider script had {} unconsumed turn(s); did the code under test make fewer LLM calls than expected?",
273 remaining.len()
274 );
275 }
276}
277
278fn take_next_turn(request: &LlmRequestPayload) -> Result<FakeLlmTurn, VmError> {
281 FAKE_LLM_CALLS.with(|calls| {
282 calls.borrow_mut().push(FakeLlmCall::from_request(request));
283 });
284 FAKE_LLM_TURNS.with(|turns| {
285 let mut turns = turns.borrow_mut();
286 if turns.is_empty() {
287 Err(VmError::Runtime(
288 "FakeLlmProvider: no script installed (or script exhausted) — install_fake_llm_script() must precede llm_call(provider: \"fake\")".to_string()
289 ))
290 } else {
291 Ok(turns.remove(0))
292 }
293 })
294}
295
296fn fake_error_to_vm_error(err: &FakeLlmError) -> VmError {
299 let message = match err.retry_after_ms {
300 Some(ms) => {
301 let secs = (ms as f64 / 1000.0).max(0.0);
302 let sep = if err.message.is_empty() || err.message.ends_with('\n') {
303 ""
304 } else {
305 "\n"
306 };
307 format!("{}{sep}retry-after: {secs}\n", err.message)
308 }
309 None => err.message.clone(),
310 };
311 VmError::CategorizedError {
312 message,
313 category: err.category.clone(),
314 }
315}
316
317pub(crate) struct FakeLlmProvider;
319
320impl LlmProvider for FakeLlmProvider {
321 fn name(&self) -> &'static str {
322 "fake"
323 }
324
325 fn requires_model(&self) -> bool {
326 false
327 }
328}
329
330impl LlmProviderChat for FakeLlmProvider {
331 fn chat<'a>(
332 &'a self,
333 request: &'a LlmRequestPayload,
334 delta_tx: Option<DeltaSender>,
335 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<LlmResult, VmError>> + 'a>> {
336 Box::pin(self.chat_impl(request, delta_tx))
337 }
338}
339
340impl FakeLlmProvider {
341 pub(crate) fn should_intercept(provider: &str) -> bool {
343 provider == "fake"
344 }
345
346 pub(crate) async fn chat_impl(
347 &self,
348 request: &LlmRequestPayload,
349 delta_tx: Option<DeltaSender>,
350 ) -> Result<LlmResult, VmError> {
351 loop {
352 let turn = take_next_turn(request)?;
353 match turn {
354 FakeLlmTurn::Stalled(duration) => {
355 if !duration.is_zero() {
356 tokio::time::sleep(duration).await;
357 }
358 }
362 FakeLlmTurn::Error(err) => {
363 return Err(fake_error_to_vm_error(&err));
364 }
365 FakeLlmTurn::Stream(events) => {
366 return play_stream(request, events, delta_tx).await;
367 }
368 }
369 }
370 }
371}
372
373async fn play_stream(
374 request: &LlmRequestPayload,
375 events: Vec<FakeLlmEvent>,
376 delta_tx: Option<DeltaSender>,
377) -> Result<LlmResult, VmError> {
378 let mut text = String::new();
379 let mut tool_calls: Vec<serde_json::Value> = Vec::new();
380 let mut blocks: Vec<serde_json::Value> = Vec::new();
381 let mut stop_reason: Option<FakeStopReason> = None;
382 let mut next_tool_index: usize = 1;
383 let mut schema_watch = crate::llm::api::StreamSchemaWatch::from_payload(request);
387
388 for event in events {
389 match event {
390 FakeLlmEvent::Token(chunk) => {
391 if let Some(tx) = delta_tx.as_ref() {
392 let _ = tx.send(chunk.clone());
393 }
394 text.push_str(&chunk);
395 if let Some(watch) = schema_watch.as_mut() {
396 if let Some(abort) = watch.observe(&chunk) {
397 return Err(abort.into_vm_error());
398 }
399 }
400 }
401 FakeLlmEvent::ToolCallDelta {
402 id,
403 name,
404 arguments,
405 } => {
406 let id = if id.is_empty() {
407 let auto = format!("fake_call_{next_tool_index}");
408 next_tool_index += 1;
409 auto
410 } else {
411 id
412 };
413 tool_calls.push(serde_json::json!({
414 "id": id,
415 "type": "tool_call",
416 "name": name,
417 "arguments": arguments,
418 }));
419 blocks.push(serde_json::json!({
420 "type": "tool_call",
421 "id": id,
422 "name": name,
423 "arguments": arguments,
424 "visibility": "internal",
425 }));
426 }
427 FakeLlmEvent::Stall(duration) => {
428 if !duration.is_zero() {
429 tokio::time::sleep(duration).await;
430 }
431 }
432 FakeLlmEvent::Error(error) => return Err(fake_error_to_vm_error(&error)),
433 FakeLlmEvent::Done(reason) => {
434 stop_reason = Some(reason);
435 break;
436 }
437 }
438 }
439
440 if !text.is_empty() {
441 let text_block = serde_json::json!({
444 "type": "output_text",
445 "text": text,
446 "visibility": "public",
447 });
448 blocks.insert(0, text_block);
449 }
450
451 let stop_reason = stop_reason.unwrap_or(if tool_calls.is_empty() {
452 FakeStopReason::EndTurn
453 } else {
454 FakeStopReason::ToolUse
455 });
456
457 Ok(LlmResult {
458 served_fast: false,
459 text,
460 raw_tool_calls: Vec::new(),
461 tool_calls,
462 input_tokens: count_input_tokens(&request.messages),
463 output_tokens: 0,
464 cache_read_tokens: 0,
465 cache_write_tokens: 0,
466 cache_supported: true,
467 model: request.model.clone(),
468 provider: "fake".to_string(),
469 thinking: None,
470 thinking_summary: None,
471 stop_reason: Some(stop_reason.as_str().to_string()),
472 blocks,
473 logprobs: Vec::new(),
474 telemetry: ProviderTelemetry::default(),
475 })
476}
477
478fn count_input_tokens(messages: &[serde_json::Value]) -> i64 {
482 fn collect(value: &serde_json::Value, out: &mut String) {
483 match value {
484 serde_json::Value::String(text) => {
485 out.push_str(text);
486 out.push('\n');
487 }
488 serde_json::Value::Array(items) => {
489 for item in items {
490 collect(item, out);
491 }
492 }
493 serde_json::Value::Object(map) => {
494 for value in map.values() {
495 collect(value, out);
496 }
497 }
498 _ => {}
499 }
500 }
501 let mut buf = String::new();
502 for message in messages {
503 collect(message, &mut buf);
504 }
505 buf.len() as i64
506}
507
508#[cfg(test)]
509mod tests {
510 use super::*;
511 use crate::llm::api::{LlmApiMode, ThinkingConfig};
512 use crate::llm::api::{LlmRequestPayload, OutputFormat};
513
514 fn fake_request() -> LlmRequestPayload {
515 LlmRequestPayload {
516 provider: "fake".to_string(),
517 model: "fake-model".to_string(),
518 region: None,
519 api_key: String::new(),
520 api_mode: LlmApiMode::ChatCompletions,
521 messages: vec![serde_json::json!({"role": "user", "content": "hello"})],
522 system: None,
523 max_tokens: 64,
524 temperature: None,
525 top_p: None,
526 top_k: None,
527 logprobs: false,
528 top_logprobs: None,
529 stop: None,
530 seed: None,
531 frequency_penalty: None,
532 presence_penalty: None,
533 fast: false,
534 output_format: OutputFormat::Text,
535 response_format: None,
536 json_schema: None,
537 output_schema: None,
538 schema_stream_abort: false,
539 thinking: ThinkingConfig::Disabled,
540 anthropic_beta_features: Vec::new(),
541 vision: false,
542 native_tools: None,
543 provider_tools: Vec::new(),
544 tool_choice: None,
545 cache: false,
546 prompt_cache_ttl: None,
547 timeout: None,
548 stream: true,
549 provider_overrides: None,
550 previous_response_id: None,
551 store: None,
552 background: None,
553 truncation: None,
554 compact: None,
555 include: None,
556 max_tool_calls: None,
557 prefill: None,
558 session_id: None,
559 reminder_lifecycle: Vec::new(),
560 cli_llm_mock_scope: None,
561 }
562 }
563
564 fn current_thread_runtime() -> tokio::runtime::Runtime {
565 tokio::runtime::Builder::new_current_thread()
566 .enable_all()
567 .start_paused(false)
568 .build()
569 .expect("runtime")
570 }
571
572 #[test]
573 fn streaming_turn_emits_deltas_in_order() {
574 let runtime = current_thread_runtime();
575 let _guard = install_fake_llm_script(FakeLlmScript::streaming(vec![
576 FakeLlmEvent::Token("hello ".into()),
577 FakeLlmEvent::Token("world".into()),
578 FakeLlmEvent::Done(FakeStopReason::EndTurn),
579 ]));
580
581 runtime.block_on(async {
582 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();
583 let result = FakeLlmProvider
584 .chat_impl(&fake_request(), Some(tx))
585 .await
586 .expect("fake call should succeed");
587
588 let mut deltas = Vec::new();
589 while let Ok(delta) = rx.try_recv() {
590 deltas.push(delta);
591 }
592
593 assert_eq!(deltas, vec!["hello ".to_string(), "world".to_string()]);
594 assert_eq!(result.text, "hello world");
595 assert_eq!(result.provider, "fake");
596 assert_eq!(result.stop_reason.as_deref(), Some("end_turn"));
597 assert_eq!(result.blocks.len(), 1);
598 assert_eq!(result.blocks[0]["type"].as_str(), Some("output_text"));
599 assert_eq!(result.blocks[0]["text"].as_str(), Some("hello world"));
600 });
601 assert_eq!(fake_llm_captured_calls().len(), 1);
602 }
603
604 #[test]
605 fn tool_call_deltas_become_tool_calls_and_blocks() {
606 let runtime = current_thread_runtime();
607 let _guard = install_fake_llm_script(FakeLlmScript::streaming(vec![
608 FakeLlmEvent::Token("calling tool".into()),
609 FakeLlmEvent::ToolCallDelta {
610 id: String::new(),
611 name: "search".into(),
612 arguments: serde_json::json!({"q": "harn"}),
613 },
614 FakeLlmEvent::Done(FakeStopReason::ToolUse),
615 ]));
616
617 runtime.block_on(async {
618 let result = FakeLlmProvider
619 .chat_impl(&fake_request(), None)
620 .await
621 .expect("fake call should succeed");
622
623 assert_eq!(result.tool_calls.len(), 1);
624 assert_eq!(result.tool_calls[0]["name"].as_str(), Some("search"));
625 assert_eq!(result.tool_calls[0]["id"].as_str(), Some("fake_call_1"));
626 assert_eq!(
627 result.tool_calls[0]["arguments"]["q"].as_str(),
628 Some("harn")
629 );
630 assert_eq!(result.stop_reason.as_deref(), Some("tool_use"));
631 assert_eq!(result.blocks[0]["type"].as_str(), Some("output_text"));
633 assert_eq!(result.blocks[1]["type"].as_str(), Some("tool_call"));
634 assert_eq!(result.blocks[1]["name"].as_str(), Some("search"));
635 });
636 }
637
638 #[test]
639 fn error_turn_returns_categorized_error() {
640 let runtime = current_thread_runtime();
641 let _guard = install_fake_llm_script(FakeLlmScript::erroring(
642 ErrorCategory::RateLimit,
643 "throttled",
644 ));
645
646 runtime.block_on(async {
647 let err = FakeLlmProvider
648 .chat_impl(&fake_request(), None)
649 .await
650 .expect_err("fake error turn should fail");
651 match err {
652 VmError::CategorizedError { message, category } => {
653 assert_eq!(category, ErrorCategory::RateLimit);
654 assert!(
655 message.contains("throttled"),
656 "error message should pass through: {message}"
657 );
658 }
659 other => panic!("expected CategorizedError, got {other:?}"),
660 }
661 });
662 }
663
664 #[test]
665 fn error_turn_embeds_retry_after_hint() {
666 let runtime = current_thread_runtime();
667 let _guard = install_fake_llm_script(FakeLlmScript::default().push(FakeLlmTurn::Error(
668 FakeLlmError::new(ErrorCategory::RateLimit, "throttled").with_retry_after_ms(2_500),
669 )));
670
671 runtime.block_on(async {
672 let err = FakeLlmProvider
673 .chat_impl(&fake_request(), None)
674 .await
675 .expect_err("fake error turn should fail");
676 let VmError::CategorizedError { message, .. } = err else {
677 panic!("expected CategorizedError");
678 };
679 assert!(
680 message.contains("retry-after: 2.5"),
681 "retry-after hint should be present in synthetic message: {message}"
682 );
683 });
684 }
685
686 #[test]
687 fn stalled_turn_advances_under_paused_clock() {
688 let runtime = tokio::runtime::Builder::new_current_thread()
689 .enable_all()
690 .start_paused(true)
691 .build()
692 .expect("paused runtime");
693 let _guard = install_fake_llm_script(
694 FakeLlmScript::default()
695 .push(FakeLlmTurn::Stalled(Duration::from_mins(1)))
696 .push(FakeLlmTurn::stream(vec![
697 FakeLlmEvent::Token("done".into()),
698 FakeLlmEvent::Done(FakeStopReason::EndTurn),
699 ])),
700 );
701
702 runtime.block_on(async {
703 let request = fake_request();
704 let chat = FakeLlmProvider.chat_impl(&request, None);
705 tokio::pin!(chat);
706
707 let polled = futures::poll!(&mut chat);
711 assert!(
712 matches!(polled, std::task::Poll::Pending),
713 "fake provider should be parked on the stall"
714 );
715
716 tokio::time::advance(Duration::from_mins(1)).await;
717 let result = chat.await.expect("after advance, fake call resolves");
718 assert_eq!(result.text, "done");
719 });
720 }
721
722 #[test]
723 fn multiple_turns_consumed_in_fifo_order() {
724 let runtime = current_thread_runtime();
725 let _guard = install_fake_llm_script(
726 FakeLlmScript::default()
727 .push(FakeLlmTurn::stream(vec![
728 FakeLlmEvent::Token("first".into()),
729 FakeLlmEvent::Done(FakeStopReason::EndTurn),
730 ]))
731 .push(FakeLlmTurn::stream(vec![
732 FakeLlmEvent::Token("second".into()),
733 FakeLlmEvent::Done(FakeStopReason::EndTurn),
734 ])),
735 );
736
737 runtime.block_on(async {
738 let first = FakeLlmProvider
739 .chat_impl(&fake_request(), None)
740 .await
741 .expect("first call");
742 let second = FakeLlmProvider
743 .chat_impl(&fake_request(), None)
744 .await
745 .expect("second call");
746 assert_eq!(first.text, "first");
747 assert_eq!(second.text, "second");
748 });
749
750 let calls = fake_llm_captured_calls();
751 assert_eq!(calls.len(), 2);
752 assert!(calls.iter().all(|c| c.provider == "fake"));
753 }
754
755 #[test]
756 #[should_panic(expected = "no script installed")]
757 fn calling_without_script_panics_with_explanatory_error() {
758 let runtime = current_thread_runtime();
759 runtime
761 .block_on(async {
762 FakeLlmProvider
763 .chat_impl(&fake_request(), None)
764 .await
765 .map_err(|e| e.to_string())
766 })
767 .unwrap();
768 }
769
770 #[test]
771 #[should_panic(expected = "unconsumed turn")]
772 fn drop_guard_asserts_on_unused_turns() {
773 let guard =
774 install_fake_llm_script(FakeLlmScript::default().push(FakeLlmTurn::stream(vec![
775 FakeLlmEvent::Done(FakeStopReason::EndTurn),
776 ])));
777 drop(guard);
779 }
780}