everruns_core/
llmsim_driver.rs

1// LLM Simulator Driver
2//
3// This module provides a fake LLM driver for testing purposes using llmsim.
4// It supports:
5// - Configurable response generators (fixed, lorem, echo, sequence)
6// - Optional tool call responses
7// - Configurable latency simulation
8// - Token counting
9//
10// Design: This driver is intended for unit and integration tests.
11// It can be configured per-test to return specific responses or tool calls.
12
13use async_trait::async_trait;
14use futures::StreamExt;
15use futures::stream;
16use std::sync::Arc;
17use std::sync::atomic::{AtomicUsize, Ordering};
18
19use crate::driver_registry::{
20    BoxedChatDriver, ChatDriver, DriverId, DriverRegistry, LlmCallConfig, LlmCompletionMetadata,
21    LlmMessage, LlmMessageRole, LlmResponseStream, LlmStreamEvent,
22};
23use crate::error::{AgentLoopError, Result};
24use crate::tool_types::ToolCall;
25use llmsim::generator::{LoremGenerator, ResponseGenerator};
26use llmsim::latency::LatencyProfile;
27use llmsim::openai::{ChatCompletionRequest, Message, Role, Usage};
28use llmsim::script::auto_tool_call_id;
29use llmsim::stream::TokenStreamBuilder;
30
31// ============================================================================
32// Configuration Types
33// ============================================================================
34
35/// Configuration for the LlmSim driver
36#[derive(Debug, Clone)]
37pub struct LlmSimConfig {
38    /// Response generation configuration
39    pub response: ResponseConfig,
40    /// Optional tool calls to include in responses
41    pub tool_calls: Option<ToolCallConfig>,
42    /// Enable latency simulation (default: false for fast tests)
43    pub simulate_latency: bool,
44    /// Model name to report in metadata
45    pub model_name: String,
46    /// Optional delay before responding (TTFT - time to first token).
47    /// This is useful for testing cancellation scenarios where we need a
48    /// predictable time window to cancel an active turn before completion.
49    pub response_delay: Option<std::time::Duration>,
50    /// Optional response ID to include in completion metadata.
51    /// Enables testing `previous_response_id` chaining.
52    pub response_id: Option<String>,
53    /// Optional capture sink for the per-call `reasoning_effort` (EVE-595).
54    /// When set, every `chat_completion_stream` call appends the effort it saw
55    /// in `LlmCallConfig`, in call order. Tests use this to assert that a
56    /// mid-turn effort change is observed by subsequent LLM steps.
57    pub effort_capture: Option<Arc<std::sync::Mutex<Vec<Option<String>>>>>,
58}
59
60impl Default for LlmSimConfig {
61    fn default() -> Self {
62        Self {
63            response: ResponseConfig::Fixed("Hello! I'm a simulated LLM response.".to_string()),
64            tool_calls: None,
65            simulate_latency: false,
66            model_name: "llmsim-model".to_string(),
67            response_delay: None,
68            response_id: None,
69            effort_capture: None,
70        }
71    }
72}
73
74impl LlmSimConfig {
75    /// Create a new config with a fixed response
76    pub fn fixed(response: impl Into<String>) -> Self {
77        Self {
78            response: ResponseConfig::Fixed(response.into()),
79            ..Default::default()
80        }
81    }
82
83    /// Create a new config that echoes user input
84    pub fn echo() -> Self {
85        Self {
86            response: ResponseConfig::Echo,
87            ..Default::default()
88        }
89    }
90
91    /// Create a new config with lorem ipsum text
92    pub fn lorem(target_tokens: usize) -> Self {
93        Self {
94            response: ResponseConfig::Lorem { target_tokens },
95            ..Default::default()
96        }
97    }
98
99    /// Create a new config with a sequence of responses
100    pub fn sequence(responses: Vec<String>) -> Self {
101        Self {
102            response: ResponseConfig::Sequence(responses),
103            ..Default::default()
104        }
105    }
106
107    /// Create a new config that replays scripted turns in order.
108    pub fn scripted(turns: Vec<SimTurn>) -> Self {
109        Self {
110            response: ResponseConfig::Scripted {
111                turns,
112                on_exhausted: OnExhausted::default(),
113            },
114            ..Default::default()
115        }
116    }
117
118    /// Set the behavior when a scripted response config exhausts its turns.
119    pub fn with_on_exhausted(mut self, mode: OnExhausted) -> Self {
120        if let ResponseConfig::Scripted { on_exhausted, .. } = &mut self.response {
121            *on_exhausted = mode;
122        }
123        self
124    }
125
126    /// Add tool calls to the response
127    pub fn with_tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
128        self.tool_calls = Some(ToolCallConfig::Fixed(tool_calls));
129        self
130    }
131
132    /// Add a sequence of tool calls (different per call)
133    pub fn with_tool_call_sequence(mut self, sequences: Vec<Vec<ToolCall>>) -> Self {
134        self.tool_calls = Some(ToolCallConfig::Sequence(sequences));
135        self
136    }
137
138    /// Enable latency simulation
139    pub fn with_latency(mut self) -> Self {
140        self.simulate_latency = true;
141        self
142    }
143
144    /// Set model name for metadata
145    pub fn with_model(mut self, model: impl Into<String>) -> Self {
146        self.model_name = model.into();
147        self
148    }
149
150    /// Set a delay before responding (TTFT - time to first token).
151    /// This creates a predictable time window for testing cancellation scenarios.
152    pub fn with_response_delay(mut self, delay: std::time::Duration) -> Self {
153        self.response_delay = Some(delay);
154        self
155    }
156
157    /// Set a response ID to include in completion metadata (for testing chaining)
158    pub fn with_response_id(mut self, id: impl Into<String>) -> Self {
159        self.response_id = Some(id.into());
160        self
161    }
162
163    /// Set a shared capture sink for the per-call `reasoning_effort` (EVE-595).
164    /// Every `chat_completion_stream` call appends the effort it observed in
165    /// `LlmCallConfig`, in call order.
166    pub fn with_effort_capture(
167        mut self,
168        capture: Arc<std::sync::Mutex<Vec<Option<String>>>>,
169    ) -> Self {
170        self.effort_capture = Some(capture);
171        self
172    }
173
174    /// Create a new config that returns an error (for testing error handling)
175    pub fn error(message: impl Into<String>) -> Self {
176        Self {
177            response: ResponseConfig::Error(message.into()),
178            ..Default::default()
179        }
180    }
181
182    /// Create a new config that returns a model-not-available error
183    pub fn model_not_available() -> Self {
184        Self {
185            response: ResponseConfig::ModelNotAvailable,
186            ..Default::default()
187        }
188    }
189}
190
191/// Response generation configuration
192#[derive(Debug, Clone)]
193pub enum ResponseConfig {
194    /// Return a fixed response
195    Fixed(String),
196    /// Echo back the last user message with a prefix
197    Echo,
198    /// Generate lorem ipsum text with target token count
199    Lorem { target_tokens: usize },
200    /// Return responses from a sequence (cycles when exhausted)
201    Sequence(Vec<String>),
202    /// Replay scripted assistant turns for multi-turn agent scenario tests.
203    Scripted {
204        turns: Vec<SimTurn>,
205        on_exhausted: OnExhausted,
206    },
207    /// Empty response (useful for tool-only responses)
208    Empty,
209    /// Simulate an error (useful for testing error handling)
210    Error(String),
211    /// Simulate a model-not-available error
212    ModelNotAvailable,
213}
214
215/// A single scripted assistant turn.
216#[derive(Debug, Clone, PartialEq)]
217pub enum SimTurn {
218    /// Plain assistant text response.
219    Assistant(String),
220    /// One or more tool calls in a single assistant turn.
221    ToolCalls(Vec<SimToolCall>),
222    /// Mixed assistant text and tool calls in the same turn.
223    Mixed {
224        text: String,
225        tool_calls: Vec<SimToolCall>,
226    },
227    /// Simulate an API/transport error on this turn.
228    Error(SimError),
229}
230
231/// A single tool call inside a scripted turn.
232#[derive(Debug, Clone, PartialEq)]
233pub struct SimToolCall {
234    pub name: String,
235    pub arguments: serde_json::Value,
236    pub id: Option<String>,
237}
238
239/// Error to inject for a scripted turn.
240#[derive(Debug, Clone, PartialEq)]
241pub enum SimError {
242    RateLimit,
243    Timeout,
244    InvalidResponse(String),
245    Other(String),
246}
247
248impl SimError {
249    fn status_code(&self) -> u16 {
250        match self {
251            SimError::RateLimit => 429,
252            SimError::Timeout => 504,
253            SimError::InvalidResponse(_) => 400,
254            SimError::Other(_) => 500,
255        }
256    }
257
258    fn message(&self) -> String {
259        match self {
260            SimError::RateLimit => "Rate limit exceeded. Please retry after some time.".to_string(),
261            SimError::Timeout => "Request timed out".to_string(),
262            SimError::InvalidResponse(message) | SimError::Other(message) => message.clone(),
263        }
264    }
265}
266
267/// Behavior when a scripted config has consumed all turns.
268#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
269pub enum OnExhausted {
270    /// Repeat the last turn forever.
271    #[default]
272    RepeatLast,
273    /// Return an error when the script is exhausted.
274    Error,
275    /// Cycle through the script from the start.
276    Loop,
277}
278
279/// Tool call configuration
280#[derive(Debug, Clone)]
281pub enum ToolCallConfig {
282    /// Always return these tool calls
283    Fixed(Vec<ToolCall>),
284    /// Return tool calls from a sequence (cycles when exhausted)
285    Sequence(Vec<Vec<ToolCall>>),
286    /// Conditionally return tool calls based on message content
287    Conditional {
288        /// Patterns to match against user message
289        patterns: Vec<ToolCallPattern>,
290    },
291}
292
293/// Pattern for conditional tool calls
294#[derive(Debug, Clone)]
295pub struct ToolCallPattern {
296    /// Substring to match in user message
297    pub contains: String,
298    /// Tool calls to return when pattern matches
299    pub tool_calls: Vec<ToolCall>,
300}
301
302impl ToolCallPattern {
303    pub fn new(contains: impl Into<String>, tool_calls: Vec<ToolCall>) -> Self {
304        Self {
305            contains: contains.into(),
306            tool_calls,
307        }
308    }
309}
310
311fn materialize_scripted_tool_calls(
312    turn_index: usize,
313    calls: Vec<SimToolCall>,
314) -> Option<Vec<ToolCall>> {
315    if calls.is_empty() {
316        return None;
317    }
318
319    Some(
320        calls
321            .into_iter()
322            .enumerate()
323            .map(|(call_index, call)| ToolCall {
324                id: call
325                    .id
326                    .unwrap_or_else(|| auto_tool_call_id(turn_index, call_index)),
327                name: call.name,
328                arguments: call.arguments,
329            })
330            .collect(),
331    )
332}
333
334// ============================================================================
335// Driver Implementation
336// ============================================================================
337
338/// LLM Simulator Driver for testing
339///
340/// This driver generates simulated responses based on configuration.
341/// It's intended for unit and integration tests where you need
342/// deterministic or configurable LLM behavior.
343///
344/// # Example
345///
346/// ```ignore
347/// use everruns_core::llmsim_driver::{LlmSimDriver, LlmSimConfig};
348///
349/// // Simple fixed response
350/// let driver = LlmSimDriver::new(LlmSimConfig::fixed("Hello!"));
351///
352/// // With tool calls
353/// let driver = LlmSimDriver::new(
354///     LlmSimConfig::fixed("Let me check that for you.")
355///         .with_tool_calls(vec![ToolCall { ... }])
356/// );
357///
358/// // Sequence of responses for multi-turn tests
359/// let driver = LlmSimDriver::new(
360///     LlmSimConfig::sequence(vec![
361///         "First response".to_string(),
362///         "Second response".to_string(),
363///     ])
364/// );
365/// ```
366#[derive(Clone)]
367pub struct LlmSimDriver {
368    config: LlmSimConfig,
369    /// Counter for sequence-based responses
370    response_counter: Arc<AtomicUsize>,
371    /// Counter for sequence-based tool calls
372    tool_call_counter: Arc<AtomicUsize>,
373}
374
375struct GeneratedTurn {
376    text: String,
377    tool_calls: Option<Vec<ToolCall>>,
378}
379
380impl LlmSimDriver {
381    /// Create a new driver with the given configuration
382    pub fn new(config: LlmSimConfig) -> Self {
383        Self {
384            config,
385            response_counter: Arc::new(AtomicUsize::new(0)),
386            tool_call_counter: Arc::new(AtomicUsize::new(0)),
387        }
388    }
389
390    /// Create a driver with default configuration (fixed response)
391    pub fn default_driver() -> Self {
392        Self::new(LlmSimConfig::default())
393    }
394
395    /// Generate response text based on configuration
396    fn generate_response(&self, messages: &[LlmMessage]) -> String {
397        match &self.config.response {
398            ResponseConfig::Fixed(text) => text.clone(),
399
400            ResponseConfig::Echo => {
401                // Find last user message and echo it
402                let last_user = messages
403                    .iter()
404                    .rev()
405                    .find(|m| m.role == LlmMessageRole::User)
406                    .map(|m| m.content_as_text())
407                    .unwrap_or_default();
408                format!("Echo: {}", last_user)
409            }
410
411            ResponseConfig::Lorem { target_tokens } => {
412                let generator = LoremGenerator::new(*target_tokens);
413                let request = self.to_chat_request(messages);
414                generator.generate(&request)
415            }
416
417            ResponseConfig::Sequence(responses) => {
418                if responses.is_empty() {
419                    return String::new();
420                }
421                let idx = self.response_counter.fetch_add(1, Ordering::SeqCst);
422                responses[idx % responses.len()].clone()
423            }
424
425            ResponseConfig::Empty => String::new(),
426
427            // Error/ModelNotAvailable cases should never be reached; checked in chat_completion_stream
428            ResponseConfig::Error(_)
429            | ResponseConfig::ModelNotAvailable
430            | ResponseConfig::Scripted { .. } => {
431                unreachable!("Special configs handled in chat_completion_stream")
432            }
433        }
434    }
435
436    /// Get tool calls based on configuration
437    fn get_tool_calls(&self, messages: &[LlmMessage]) -> Option<Vec<ToolCall>> {
438        match &self.config.tool_calls {
439            None => None,
440
441            Some(ToolCallConfig::Fixed(calls)) => {
442                if calls.is_empty() {
443                    None
444                } else {
445                    Some(calls.clone())
446                }
447            }
448
449            Some(ToolCallConfig::Sequence(sequences)) => {
450                if sequences.is_empty() {
451                    return None;
452                }
453                let idx = self.tool_call_counter.fetch_add(1, Ordering::SeqCst);
454                let calls = &sequences[idx % sequences.len()];
455                if calls.is_empty() {
456                    None
457                } else {
458                    Some(calls.clone())
459                }
460            }
461
462            Some(ToolCallConfig::Conditional { patterns }) => {
463                // Scan user messages newest-first and use the first one that
464                // matches a pattern. Looking only at the very last user message
465                // is brittle: an injected user-role notification (e.g. a
466                // background task's terminal wake-up) can land after the
467                // triggering prompt and mask it, even though content-keyed
468                // patterns are meant to make scheduling order irrelevant.
469                // Newest-first honours the most recent matching intent while
470                // skipping interleaved non-matching notifications.
471                for message in messages.iter().rev() {
472                    if message.role != LlmMessageRole::User {
473                        continue;
474                    }
475                    let text = message.content_as_text();
476                    if let Some(pattern) = patterns.iter().find(|p| text.contains(&p.contains)) {
477                        return if pattern.tool_calls.is_empty() {
478                            None
479                        } else {
480                            Some(pattern.tool_calls.clone())
481                        };
482                    }
483                }
484                None
485            }
486        }
487    }
488
489    fn generate_turn(&self, messages: &[LlmMessage]) -> Result<GeneratedTurn> {
490        if let ResponseConfig::Scripted {
491            turns,
492            on_exhausted,
493        } = &self.config.response
494        {
495            return self.generate_scripted_turn(turns, *on_exhausted);
496        }
497
498        Ok(GeneratedTurn {
499            text: self.generate_response(messages),
500            tool_calls: self.get_tool_calls(messages),
501        })
502    }
503
504    fn generate_scripted_turn(
505        &self,
506        turns: &[SimTurn],
507        on_exhausted: OnExhausted,
508    ) -> Result<GeneratedTurn> {
509        if turns.is_empty() {
510            return Err(AgentLoopError::config(
511                "llmsim scripted config must contain at least one turn",
512            ));
513        }
514
515        let turn_index = self.response_counter.fetch_add(1, Ordering::SeqCst);
516        let turn = if turn_index < turns.len() {
517            turns[turn_index].clone()
518        } else {
519            match on_exhausted {
520                OnExhausted::RepeatLast => turns[turns.len() - 1].clone(),
521                OnExhausted::Loop => turns[turn_index % turns.len()].clone(),
522                OnExhausted::Error => {
523                    return Err(AgentLoopError::config("llmsim scripted config exhausted"));
524                }
525            }
526        };
527
528        match turn {
529            SimTurn::Assistant(text) => Ok(GeneratedTurn {
530                text,
531                tool_calls: None,
532            }),
533            SimTurn::ToolCalls(calls) => Ok(GeneratedTurn {
534                text: String::new(),
535                tool_calls: materialize_scripted_tool_calls(turn_index, calls),
536            }),
537            SimTurn::Mixed { text, tool_calls } => Ok(GeneratedTurn {
538                text,
539                tool_calls: materialize_scripted_tool_calls(turn_index, tool_calls),
540            }),
541            SimTurn::Error(error) => Err(AgentLoopError::llm(format!(
542                "LlmSim scripted error ({}): {}",
543                error.status_code(),
544                error.message()
545            ))),
546        }
547    }
548
549    /// Convert LlmMessage to llmsim ChatCompletionRequest
550    fn to_chat_request(&self, messages: &[LlmMessage]) -> ChatCompletionRequest {
551        let sim_messages: Vec<Message> = messages
552            .iter()
553            .map(|m| {
554                let role = match m.role {
555                    LlmMessageRole::System => Role::System,
556                    LlmMessageRole::User => Role::User,
557                    LlmMessageRole::Assistant => Role::Assistant,
558                    LlmMessageRole::Tool => Role::Tool,
559                };
560                Message {
561                    role,
562                    content: Some(m.content_as_text()),
563                    name: None,
564                    tool_calls: None,
565                    tool_call_id: m.tool_call_id.clone(),
566                }
567            })
568            .collect();
569
570        ChatCompletionRequest {
571            model: self.config.model_name.clone(),
572            messages: sim_messages,
573            temperature: None,
574            top_p: None,
575            n: None,
576            max_tokens: None,
577            max_completion_tokens: None,
578            stream: true,
579            stop: None,
580            presence_penalty: None,
581            frequency_penalty: None,
582            logit_bias: None,
583            user: None,
584            tools: None,
585            tool_choice: None,
586            seed: None,
587            response_format: None,
588        }
589    }
590
591    /// Resolve latency profile for a request.
592    /// Model names containing "-latency" enable realistic streaming simulation
593    /// via LatencyProfile::fast(). The config flag `simulate_latency` also enables it.
594    /// Returns LatencyProfile::instant() when neither is set.
595    fn resolve_latency_profile(&self, model_name: &str) -> LatencyProfile {
596        if self.config.simulate_latency || model_name.contains("-latency") {
597            LatencyProfile::fast()
598        } else {
599            LatencyProfile::instant()
600        }
601    }
602
603    /// Estimate token count for text
604    fn estimate_tokens(text: &str) -> u32 {
605        // Simple estimation: ~4 chars per token
606        (text.len() / 4).max(1) as u32
607    }
608}
609
610#[async_trait]
611impl ChatDriver for LlmSimDriver {
612    async fn chat_completion_stream(
613        &self,
614        messages: Vec<LlmMessage>,
615        config: &LlmCallConfig,
616    ) -> Result<LlmResponseStream> {
617        // Record the per-call reasoning effort for tests (EVE-595). Captured
618        // before any error short-circuit so even error turns are observable.
619        if let Some(capture) = &self.config.effort_capture
620            && let Ok(mut efforts) = capture.lock()
621        {
622            efforts.push(config.reasoning_effort.clone());
623        }
624
625        // Check for error configs first
626        if let ResponseConfig::Error(error_msg) = &self.config.response {
627            return Err(anyhow::anyhow!("LLM error: {}", error_msg).into());
628        }
629        if matches!(self.config.response, ResponseConfig::ModelNotAvailable) {
630            return Err(AgentLoopError::model_not_available(config.model.clone()));
631        }
632
633        // Apply response delay if configured or if model name contains "-ttft-{ms}".
634        // TTFT = Time To First Token. This simulates LLM "thinking" time.
635        // Used for testing cancellation scenarios where we need a predictable
636        // time window to cancel an active turn before the LLM completes.
637        let delay = self
638            .config
639            .response_delay
640            .or_else(|| parse_ttft_from_model_name(&config.model));
641        if let Some(delay) = delay {
642            tokio::time::sleep(delay).await;
643        }
644
645        let generated_turn = self.generate_turn(&messages)?;
646        let response_text = generated_turn.text;
647        let tool_calls = generated_turn.tool_calls;
648        let model_name = config.model.clone();
649        let response_id_for_done = self.config.response_id.clone();
650        let latency_profile = self.resolve_latency_profile(&model_name);
651
652        // Calculate token estimates
653        let prompt_tokens: u32 = messages
654            .iter()
655            .map(|m| Self::estimate_tokens(&m.content_as_text()))
656            .sum();
657        let completion_tokens = Self::estimate_tokens(&response_text);
658
659        // Use llmsim's TokenStreamBuilder for streaming with latency simulation.
660        // It handles TTFT and inter-token delays natively via LatencyProfile.
661        let usage = Usage {
662            prompt_tokens,
663            completion_tokens,
664            total_tokens: prompt_tokens + completion_tokens,
665        };
666
667        let chunk_stream = TokenStreamBuilder::new(&model_name, &response_text)
668            .latency(latency_profile)
669            .usage(usage)
670            .build()
671            .into_chunk_stream();
672
673        // Map llmsim ChatCompletionChunk -> our LlmStreamEvent, then append
674        // tool calls and metadata after the text stream completes.
675        let tool_calls_tail = tool_calls;
676        let model_name_done = model_name.clone();
677        let event_stream = chunk_stream.flat_map(move |chunk| {
678            let mut events: Vec<Result<LlmStreamEvent>> = Vec::new();
679
680            for choice in &chunk.choices {
681                if let Some(content) = &choice.delta.content
682                    && !content.is_empty()
683                {
684                    events.push(Ok(LlmStreamEvent::TextDelta(content.clone())));
685                }
686            }
687
688            stream::iter(events)
689        });
690
691        // Append tool calls + done after the text stream
692        let done_events: Vec<Result<LlmStreamEvent>> = {
693            let mut tail = Vec::new();
694            if let Some(calls) = tool_calls_tail {
695                tail.push(Ok(LlmStreamEvent::ToolCalls(calls)));
696            }
697            tail.push(Ok(LlmStreamEvent::Done(Box::new(LlmCompletionMetadata {
698                total_tokens: Some(prompt_tokens + completion_tokens),
699                prompt_tokens: Some(prompt_tokens),
700                completion_tokens: Some(completion_tokens),
701                cache_read_tokens: None,
702                cache_creation_tokens: None,
703                provider_cost_usd: None,
704                model: Some(model_name_done),
705                finish_reason: Some("stop".to_string()),
706                retry_metadata: None,
707                response_id: response_id_for_done,
708                phase: None,
709            }))));
710            tail
711        };
712
713        let full_stream = event_stream.chain(stream::iter(done_events));
714        Ok(Box::pin(full_stream))
715    }
716}
717
718impl std::fmt::Debug for LlmSimDriver {
719    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
720        f.debug_struct("LlmSimDriver")
721            .field("model", &self.config.model_name)
722            .field("simulate_latency", &self.config.simulate_latency)
723            .finish()
724    }
725}
726
727// ============================================================================
728// Driver Registration
729// ============================================================================
730
731/// Register the LlmSim driver with the driver registry
732///
733/// This registers a driver for the `LlmSim` provider type.
734/// The driver is created with a default configuration; for custom configs,
735/// create the driver directly using `LlmSimDriver::new()`.
736///
737/// # Example
738///
739/// ```ignore
740/// use everruns_core::DriverRegistry;
741/// use everruns_core::llmsim_driver::register_driver;
742///
743/// let mut registry = DriverRegistry::new();
744/// register_driver(&mut registry);
745/// ```
746pub fn register_driver(registry: &mut DriverRegistry) {
747    registry.register(DriverId::LlmSim, |_config| {
748        // Default driver - tests can create custom drivers directly
749        Box::new(LlmSimDriver::default_driver()) as BoxedChatDriver
750    });
751}
752
753/// Register the LlmSim driver with a custom configuration. Useful for
754/// servers/workers that want to opt into a scripted scenario (e.g. the
755/// `user_hooks` audit-log demo) without changing the default behaviour for
756/// callers that don't.
757///
758/// The same `config` is cloned for every constructed driver in the
759/// registry, so its `Arc`-backed counters (sequence index, etc.) are
760/// shared across invocations.
761pub fn register_driver_with_config(registry: &mut DriverRegistry, config: LlmSimConfig) {
762    let driver = LlmSimDriver::new(config);
763    registry.register(DriverId::LlmSim, move |_config| {
764        Box::new(driver.clone()) as BoxedChatDriver
765    });
766}
767
768/// Parse TTFT (time to first token) delay from model name if it contains "-ttft-{ms}" pattern.
769/// For example: "llmsim-ttft-2000" returns Some(Duration::from_millis(2000))
770///
771/// This allows tests to opt-in to response delays by using specific model names,
772/// which is useful for testing cancellation of active turns.
773fn parse_ttft_from_model_name(model_name: &str) -> Option<std::time::Duration> {
774    if let Some(idx) = model_name.find("-ttft-") {
775        let after_ttft = &model_name[idx + 6..]; // skip "-ttft-"
776        let ms_str: String = after_ttft
777            .chars()
778            .take_while(|c| c.is_ascii_digit())
779            .collect();
780        if let Ok(ms) = ms_str.parse::<u64>()
781            && ms > 0
782        {
783            return Some(std::time::Duration::from_millis(ms));
784        }
785    }
786    None
787}
788
789/// Create a LlmSim driver with custom configuration
790///
791/// This is the preferred way to create a driver in tests.
792/// Unlike `register_driver`, this gives you full control over the config.
793///
794/// # Example
795///
796/// ```ignore
797/// use everruns_core::llmsim_driver::{create_chat_driver, LlmSimConfig};
798///
799/// let driver = create_chat_driver(
800///     LlmSimConfig::fixed("I'll help you with that!")
801///         .with_tool_calls(vec![...])
802/// );
803/// ```
804pub fn create_chat_driver(config: LlmSimConfig) -> BoxedChatDriver {
805    Box::new(LlmSimDriver::new(config))
806}
807
808// ============================================================================
809// Pre-baked demo scripts
810// ============================================================================
811
812/// Scripted multi-turn config that drives the Cloud Cost & Security Auditor
813/// example agent through a small, deterministic AWS audit using only the
814/// `fake_aws` capability's tools. Useful for the `user_hooks` end-to-end
815/// demo (and any operator who wants to exercise the auditor without an LLM
816/// API key).
817///
818/// Sequence of assistant turns:
819///   1. Call `aws_list_ec2_instances`.
820///   2. Call `aws_list_s3_buckets`.
821///   3. Write a short audit summary as plain text.
822///
823/// After turn 3 the script repeats the final turn, matching the default
824/// `OnExhausted::RepeatLast` behavior.
825pub fn auditor_demo_script() -> LlmSimConfig {
826    let turns = vec![
827        SimTurn::Mixed {
828            text: "Starting the audit. Listing EC2 instances first.".to_string(),
829            tool_calls: vec![SimToolCall {
830                name: "aws_list_ec2_instances".to_string(),
831                arguments: serde_json::json!({}),
832                id: Some("call_demo_ec2".to_string()),
833            }],
834        },
835        SimTurn::Mixed {
836            text: "EC2 inventory captured. Listing S3 buckets next.".to_string(),
837            tool_calls: vec![SimToolCall {
838                name: "aws_list_s3_buckets".to_string(),
839                arguments: serde_json::json!({}),
840                id: Some("call_demo_s3".to_string()),
841            }],
842        },
843        SimTurn::Assistant(
844            "Audit complete: inventoried EC2 instances and S3 buckets. \
845             See /workspace/.audit.log for the per-tool-call audit trail \
846             written by the post_tool_use hook bundle."
847                .to_string(),
848        ),
849    ];
850    LlmSimConfig::scripted(turns)
851}
852
853/// Scripted scenario that exercises the `pre_tool_use` block path. The
854/// scripted agent first issues a destructive `bash` call (`rm -rf /`)
855/// followed by a benign one (`ls -la`). When combined with a `pre_tool_use`
856/// hook bundle that denies `rm -rf` patterns, the first tool call gets
857/// blocked (the tool is not invoked) and the second succeeds — the agent
858/// observes the difference in tool results.
859///
860/// Used by `LLMSIM_DEMO=guarded` to demonstrate `pre_tool_use` without an
861/// LLM API key.
862pub fn guarded_bash_demo_script() -> LlmSimConfig {
863    let turns = vec![
864        SimTurn::Mixed {
865            text: "Step 1: attempting a destructive command.".to_string(),
866            tool_calls: vec![SimToolCall {
867                name: "bash".to_string(),
868                arguments: serde_json::json!({ "commands": "rm -rf /" }),
869                id: Some("call_demo_rm".to_string()),
870            }],
871        },
872        SimTurn::Mixed {
873            text: "Step 2: trying a safe command.".to_string(),
874            tool_calls: vec![SimToolCall {
875                name: "bash".to_string(),
876                arguments: serde_json::json!({ "commands": "ls -la /workspace" }),
877                id: Some("call_demo_ls".to_string()),
878            }],
879        },
880        SimTurn::Assistant(
881            "Guarded-bash demo complete. The first tool call should be \
882             blocked by the pre_tool_use hook; the second should succeed."
883                .to_string(),
884        ),
885    ];
886    LlmSimConfig::scripted(turns)
887}
888
889/// Scripted scenario that exercises the session task registry end-to-end
890/// without an LLM API key. The scripted agent starts a background bash run
891/// via `spawn_background` (which creates a `background_tool` session task),
892/// then inspects the registry with `list_tasks`.
893///
894/// Used by `LLMSIM_DEMO=tasks`. Requires an agent with the `bashkit_shell`
895/// and `session_tasks` capabilities.
896pub fn session_tasks_demo_script() -> LlmSimConfig {
897    let turns = vec![
898        SimTurn::Mixed {
899            text: "Kicking off a background bash run.".to_string(),
900            tool_calls: vec![SimToolCall {
901                name: "spawn_background".to_string(),
902                arguments: serde_json::json!({
903                    "tool": "bash",
904                    "args": { "commands": "echo task demo start; echo task demo done" },
905                    "title": "Demo background run",
906                    "signal_on_completion": false,
907                }),
908                id: Some("call_demo_spawn".to_string()),
909            }],
910        },
911        SimTurn::Mixed {
912            text: "Checking the session task registry.".to_string(),
913            tool_calls: vec![SimToolCall {
914                name: "list_tasks".to_string(),
915                arguments: serde_json::json!({}),
916                id: Some("call_demo_list".to_string()),
917            }],
918        },
919        SimTurn::Assistant(
920            "Session tasks demo complete: a background run was started and \
921             tracked as a session task. Inspect it via \
922             GET /v1/sessions/{session_id}/tasks."
923                .to_string(),
924        ),
925    ];
926    LlmSimConfig::scripted(turns)
927}
928
929/// Scripted scenario for the monitor task kind: spawns a recurring scheduled
930/// monitor (cron fires at second 0 every minute) so the session scheduler
931/// creates the schedule and a linked `monitor` task. Used by
932/// `LLMSIM_DEMO=monitor` for end-to-end verification without an LLM API key.
933pub fn monitor_demo_script() -> LlmSimConfig {
934    let turns = vec![
935        SimTurn::Mixed {
936            text: "Setting up a recurring monitor.".to_string(),
937            tool_calls: vec![SimToolCall {
938                name: "spawn_background".to_string(),
939                arguments: serde_json::json!({
940                    "tool": "bash",
941                    "args": { "commands": "echo monitor check" },
942                    "title": "Demo monitor",
943                    "signal_on_completion": false,
944                    "schedule": { "cron_expression": "0 * * * * * *", "timezone": "UTC" },
945                }),
946                id: Some("call_demo_monitor".to_string()),
947            }],
948        },
949        SimTurn::Assistant(
950            "Monitor demo complete: a recurring monitor was scheduled and tracked as a session task. Inspect it via GET /v1/sessions/{session_id}/tasks."
951                .to_string(),
952        ),
953    ];
954    LlmSimConfig::scripted(turns)
955}
956
957// ============================================================================
958// Tests
959// ============================================================================
960
961#[cfg(test)]
962mod tests {
963    use super::*;
964    use futures::StreamExt;
965
966    #[test]
967    fn auditor_demo_script_calls_ec2_then_s3_then_summarises() {
968        let config = auditor_demo_script();
969        let turns = match &config.response {
970            ResponseConfig::Scripted { turns, .. } => turns,
971            other => panic!("expected Scripted, got {other:?}"),
972        };
973        assert_eq!(turns.len(), 3, "script has three turns");
974        match &turns[0] {
975            SimTurn::Mixed { tool_calls, .. } => {
976                assert_eq!(tool_calls.len(), 1);
977                assert_eq!(tool_calls[0].name, "aws_list_ec2_instances");
978            }
979            other => panic!("turn 0 should be Mixed, got {other:?}"),
980        }
981        match &turns[1] {
982            SimTurn::Mixed { tool_calls, .. } => {
983                assert_eq!(tool_calls.len(), 1);
984                assert_eq!(tool_calls[0].name, "aws_list_s3_buckets");
985            }
986            other => panic!("turn 1 should be Mixed, got {other:?}"),
987        }
988        match &turns[2] {
989            SimTurn::Assistant(text) => {
990                assert!(
991                    text.contains("/workspace/.audit.log"),
992                    "summary mentions the audit log: {text:?}"
993                );
994            }
995            other => panic!("turn 2 should be Assistant, got {other:?}"),
996        }
997    }
998
999    fn make_config() -> LlmCallConfig {
1000        LlmCallConfig {
1001            speed: None,
1002            verbosity: None,
1003            model: "test-model".to_string(),
1004            temperature: None,
1005            max_tokens: None,
1006            tools: vec![],
1007            reasoning_effort: None,
1008            metadata: std::collections::HashMap::new(),
1009            previous_response_id: None,
1010            tool_search: None,
1011            prompt_cache: None,
1012            openrouter_routing: None,
1013            parallel_tool_calls: None,
1014            volatile_suffix_len: 0,
1015        }
1016    }
1017
1018    fn user_message(content: &str) -> LlmMessage {
1019        LlmMessage::text(LlmMessageRole::User, content)
1020    }
1021
1022    fn system_message(content: &str) -> LlmMessage {
1023        LlmMessage::text(LlmMessageRole::System, content)
1024    }
1025
1026    #[tokio::test]
1027    async fn test_fixed_response() {
1028        let driver = LlmSimDriver::new(LlmSimConfig::fixed("Hello, world!"));
1029        let messages = vec![user_message("Hi there")];
1030
1031        let response = driver
1032            .chat_completion(messages, &make_config())
1033            .await
1034            .unwrap();
1035
1036        assert_eq!(response.text, "Hello, world!");
1037        assert!(response.tool_calls.is_none());
1038    }
1039
1040    #[tokio::test]
1041    async fn test_echo_response() {
1042        let driver = LlmSimDriver::new(LlmSimConfig::echo());
1043        let messages = vec![
1044            system_message("You are a helpful assistant"),
1045            user_message("What is 2+2?"),
1046        ];
1047
1048        let response = driver
1049            .chat_completion(messages, &make_config())
1050            .await
1051            .unwrap();
1052
1053        assert_eq!(response.text, "Echo: What is 2+2?");
1054    }
1055
1056    #[tokio::test]
1057    async fn test_sequence_response() {
1058        let driver = LlmSimDriver::new(LlmSimConfig::sequence(vec![
1059            "First".to_string(),
1060            "Second".to_string(),
1061            "Third".to_string(),
1062        ]));
1063
1064        let messages = vec![user_message("test")];
1065
1066        // First call
1067        let r1 = driver
1068            .chat_completion(messages.clone(), &make_config())
1069            .await
1070            .unwrap();
1071        assert_eq!(r1.text, "First");
1072
1073        // Second call
1074        let r2 = driver
1075            .chat_completion(messages.clone(), &make_config())
1076            .await
1077            .unwrap();
1078        assert_eq!(r2.text, "Second");
1079
1080        // Third call
1081        let r3 = driver
1082            .chat_completion(messages.clone(), &make_config())
1083            .await
1084            .unwrap();
1085        assert_eq!(r3.text, "Third");
1086
1087        // Fourth call - cycles back to first
1088        let r4 = driver
1089            .chat_completion(messages.clone(), &make_config())
1090            .await
1091            .unwrap();
1092        assert_eq!(r4.text, "First");
1093    }
1094
1095    #[tokio::test]
1096    async fn test_lorem_response() {
1097        let driver = LlmSimDriver::new(LlmSimConfig::lorem(50));
1098        let messages = vec![user_message("Generate text")];
1099
1100        let response = driver
1101            .chat_completion(messages, &make_config())
1102            .await
1103            .unwrap();
1104
1105        // Lorem response should have content
1106        assert!(!response.text.is_empty());
1107        // Should have multiple words
1108        assert!(response.text.split_whitespace().count() > 5);
1109    }
1110
1111    #[tokio::test]
1112    async fn test_fixed_tool_calls() {
1113        let tool_call = ToolCall {
1114            id: "call_123".to_string(),
1115            name: "get_weather".to_string(),
1116            arguments: serde_json::json!({"city": "NYC"}),
1117        };
1118
1119        let driver = LlmSimDriver::new(
1120            LlmSimConfig::fixed("Let me check the weather.")
1121                .with_tool_calls(vec![tool_call.clone()]),
1122        );
1123
1124        let messages = vec![user_message("What's the weather?")];
1125        let response = driver
1126            .chat_completion(messages, &make_config())
1127            .await
1128            .unwrap();
1129
1130        assert_eq!(response.text, "Let me check the weather.");
1131        let calls = response.tool_calls.expect("Expected tool calls");
1132        assert_eq!(calls.len(), 1);
1133        assert_eq!(calls[0].name, "get_weather");
1134        assert_eq!(calls[0].id, "call_123");
1135    }
1136
1137    #[tokio::test]
1138    async fn test_tool_call_sequence() {
1139        let call1 = ToolCall {
1140            id: "call_1".to_string(),
1141            name: "search".to_string(),
1142            arguments: serde_json::json!({"q": "rust"}),
1143        };
1144        let call2 = ToolCall {
1145            id: "call_2".to_string(),
1146            name: "fetch".to_string(),
1147            arguments: serde_json::json!({"url": "https://example.com"}),
1148        };
1149
1150        let driver = LlmSimDriver::new(
1151            LlmSimConfig::fixed("Processing...").with_tool_call_sequence(vec![
1152                vec![call1.clone()],
1153                vec![call2.clone()],
1154                vec![],
1155            ]),
1156        );
1157
1158        let messages = vec![user_message("test")];
1159
1160        // First call - should get search
1161        let r1 = driver
1162            .chat_completion(messages.clone(), &make_config())
1163            .await
1164            .unwrap();
1165        let calls1 = r1.tool_calls.expect("Expected tool calls");
1166        assert_eq!(calls1[0].name, "search");
1167
1168        // Second call - should get fetch
1169        let r2 = driver
1170            .chat_completion(messages.clone(), &make_config())
1171            .await
1172            .unwrap();
1173        let calls2 = r2.tool_calls.expect("Expected tool calls");
1174        assert_eq!(calls2[0].name, "fetch");
1175
1176        // Third call - no tool calls
1177        let r3 = driver
1178            .chat_completion(messages.clone(), &make_config())
1179            .await
1180            .unwrap();
1181        assert!(r3.tool_calls.is_none());
1182    }
1183
1184    #[tokio::test]
1185    async fn test_scripted_multi_turn_tool_call_agent_sequence() {
1186        let driver = LlmSimDriver::new(
1187            LlmSimConfig::scripted(vec![
1188                SimTurn::ToolCalls(vec![SimToolCall {
1189                    name: "bash".to_string(),
1190                    arguments: serde_json::json!({"command": "echo hello > /tmp/x.txt"}),
1191                    id: None,
1192                }]),
1193                SimTurn::ToolCalls(vec![SimToolCall {
1194                    name: "bash".to_string(),
1195                    arguments: serde_json::json!({"command": "sed -i s/hello/world/ /tmp/x.txt"}),
1196                    id: None,
1197                }]),
1198                SimTurn::Assistant("done".to_string()),
1199            ])
1200            .with_on_exhausted(OnExhausted::Error),
1201        );
1202
1203        let messages = vec![user_message("create /tmp/x.txt then change hello to world")];
1204
1205        let first = driver
1206            .chat_completion(messages.clone(), &make_config())
1207            .await
1208            .unwrap();
1209        let first_calls = first.tool_calls.expect("first turn should call bash");
1210        assert_eq!(first.text, "");
1211        assert_eq!(first_calls[0].name, "bash");
1212        assert_eq!(first_calls[0].id, "call_llmsim_0_0");
1213
1214        let second = driver
1215            .chat_completion(messages.clone(), &make_config())
1216            .await
1217            .unwrap();
1218        let second_calls = second.tool_calls.expect("second turn should call bash");
1219        assert_eq!(second_calls[0].name, "bash");
1220        assert_eq!(second_calls[0].id, "call_llmsim_1_0");
1221
1222        let final_response = driver
1223            .chat_completion(messages.clone(), &make_config())
1224            .await
1225            .unwrap();
1226        assert_eq!(final_response.text, "done");
1227        assert!(final_response.tool_calls.is_none());
1228
1229        let exhausted = driver
1230            .chat_completion(messages, &make_config())
1231            .await
1232            .unwrap_err();
1233        assert!(matches!(exhausted, AgentLoopError::Configuration(_)));
1234    }
1235
1236    #[tokio::test]
1237    async fn test_scripted_mixed_turn_streams_text_and_tool_calls() {
1238        let driver = LlmSimDriver::new(LlmSimConfig::scripted(vec![SimTurn::Mixed {
1239            text: "Let me check".to_string(),
1240            tool_calls: vec![SimToolCall {
1241                name: "search".to_string(),
1242                arguments: serde_json::json!({"q": "rust"}),
1243                id: Some("call_search".to_string()),
1244            }],
1245        }]));
1246
1247        let mut stream = driver
1248            .chat_completion_stream(vec![user_message("find rust")], &make_config())
1249            .await
1250            .unwrap();
1251
1252        let mut text_parts = Vec::new();
1253        let mut tool_calls = None;
1254        while let Some(event) = stream.next().await {
1255            match event.unwrap() {
1256                LlmStreamEvent::TextDelta(text) => text_parts.push(text),
1257                LlmStreamEvent::ToolCalls(calls) => tool_calls = Some(calls),
1258                LlmStreamEvent::Done(_) => {}
1259                _ => {}
1260            }
1261        }
1262
1263        assert!(!text_parts.is_empty(), "scripted text should stream");
1264        assert_eq!(text_parts.join(""), "Let me check");
1265        let calls = tool_calls.expect("mixed turn should emit tool calls");
1266        assert_eq!(calls[0].id, "call_search");
1267        assert_eq!(calls[0].name, "search");
1268    }
1269
1270    #[tokio::test]
1271    async fn test_scripted_on_exhausted_modes() {
1272        let repeat = LlmSimDriver::new(LlmSimConfig::scripted(vec![
1273            SimTurn::Assistant("one".to_string()),
1274            SimTurn::Assistant("two".to_string()),
1275        ]));
1276        let messages = vec![user_message("test")];
1277        assert_eq!(
1278            repeat
1279                .chat_completion(messages.clone(), &make_config())
1280                .await
1281                .unwrap()
1282                .text,
1283            "one"
1284        );
1285        assert_eq!(
1286            repeat
1287                .chat_completion(messages.clone(), &make_config())
1288                .await
1289                .unwrap()
1290                .text,
1291            "two"
1292        );
1293        assert_eq!(
1294            repeat
1295                .chat_completion(messages.clone(), &make_config())
1296                .await
1297                .unwrap()
1298                .text,
1299            "two"
1300        );
1301
1302        let looping = LlmSimDriver::new(
1303            LlmSimConfig::scripted(vec![
1304                SimTurn::Assistant("a".to_string()),
1305                SimTurn::Assistant("b".to_string()),
1306            ])
1307            .with_on_exhausted(OnExhausted::Loop),
1308        );
1309        assert_eq!(
1310            looping
1311                .chat_completion(messages.clone(), &make_config())
1312                .await
1313                .unwrap()
1314                .text,
1315            "a"
1316        );
1317        assert_eq!(
1318            looping
1319                .chat_completion(messages.clone(), &make_config())
1320                .await
1321                .unwrap()
1322                .text,
1323            "b"
1324        );
1325        assert_eq!(
1326            looping
1327                .chat_completion(messages, &make_config())
1328                .await
1329                .unwrap()
1330                .text,
1331            "a"
1332        );
1333    }
1334
1335    #[tokio::test]
1336    async fn test_scripted_error_turn() {
1337        let driver = LlmSimDriver::new(LlmSimConfig::scripted(vec![SimTurn::Error(
1338            SimError::RateLimit,
1339        )]));
1340
1341        let err = driver
1342            .chat_completion(vec![user_message("test")], &make_config())
1343            .await
1344            .unwrap_err();
1345
1346        assert!(err.is_rate_limited());
1347    }
1348
1349    #[tokio::test]
1350    async fn test_conditional_tool_calls() {
1351        let weather_call = ToolCall {
1352            id: "call_w".to_string(),
1353            name: "get_weather".to_string(),
1354            arguments: serde_json::json!({}),
1355        };
1356        let search_call = ToolCall {
1357            id: "call_s".to_string(),
1358            name: "search".to_string(),
1359            arguments: serde_json::json!({}),
1360        };
1361
1362        let config = LlmSimConfig {
1363            response: ResponseConfig::Fixed("Response".to_string()),
1364            tool_calls: Some(ToolCallConfig::Conditional {
1365                patterns: vec![
1366                    ToolCallPattern::new("weather", vec![weather_call]),
1367                    ToolCallPattern::new("search", vec![search_call]),
1368                ],
1369            }),
1370            simulate_latency: false,
1371            model_name: "test".to_string(),
1372            response_delay: None,
1373            response_id: None,
1374            effort_capture: None,
1375        };
1376
1377        let driver = LlmSimDriver::new(config);
1378
1379        // Weather query - should trigger weather tool
1380        let r1 = driver
1381            .chat_completion(vec![user_message("What's the weather?")], &make_config())
1382            .await
1383            .unwrap();
1384        let calls1 = r1.tool_calls.expect("Expected weather tool");
1385        assert_eq!(calls1[0].name, "get_weather");
1386
1387        // Search query - should trigger search tool
1388        let r2 = driver
1389            .chat_completion(vec![user_message("search for rust")], &make_config())
1390            .await
1391            .unwrap();
1392        let calls2 = r2.tool_calls.expect("Expected search tool");
1393        assert_eq!(calls2[0].name, "search");
1394
1395        // No matching pattern - no tool calls
1396        let r3 = driver
1397            .chat_completion(vec![user_message("hello world")], &make_config())
1398            .await
1399            .unwrap();
1400        assert!(r3.tool_calls.is_none());
1401    }
1402
1403    #[tokio::test]
1404    async fn test_streaming() {
1405        let driver = LlmSimDriver::new(LlmSimConfig::fixed("Hello world test"));
1406        let messages = vec![user_message("test")];
1407
1408        let mut stream = driver
1409            .chat_completion_stream(messages, &make_config())
1410            .await
1411            .unwrap();
1412
1413        let mut text_parts = Vec::new();
1414        let mut got_done = false;
1415
1416        while let Some(event) = stream.next().await {
1417            match event.unwrap() {
1418                LlmStreamEvent::TextDelta(text) => text_parts.push(text),
1419                LlmStreamEvent::Done(meta) => {
1420                    got_done = true;
1421                    assert!(meta.total_tokens.is_some());
1422                    assert!(meta.model.is_some());
1423                }
1424                _ => {}
1425            }
1426        }
1427
1428        assert!(got_done);
1429        // llmsim's TokenStream handles chunking; verify full text is correct
1430        assert!(!text_parts.is_empty());
1431        assert_eq!(text_parts.join(""), "Hello world test");
1432    }
1433
1434    #[tokio::test]
1435    async fn test_metadata() {
1436        let driver = LlmSimDriver::new(LlmSimConfig::fixed("Hi").with_model("custom-model"));
1437        let messages = vec![user_message("test")];
1438
1439        let mut config = make_config();
1440        config.model = "request-model".to_string();
1441
1442        let response = driver.chat_completion(messages, &config).await.unwrap();
1443
1444        // Model should come from the request config
1445        assert_eq!(response.metadata.model, Some("request-model".to_string()));
1446        assert!(response.metadata.prompt_tokens.is_some());
1447        assert!(response.metadata.completion_tokens.is_some());
1448    }
1449
1450    #[tokio::test]
1451    async fn test_register_driver() {
1452        let mut registry = DriverRegistry::new();
1453        register_driver(&mut registry);
1454
1455        assert!(registry.has_driver(&DriverId::LlmSim));
1456
1457        // Creating a driver should work (with any API key since it's simulated)
1458        let config =
1459            crate::driver_registry::ProviderConfig::new(DriverId::LlmSim).with_api_key("fake-key");
1460        let driver = registry.create_chat_driver(&config);
1461        assert!(driver.is_ok());
1462    }
1463
1464    #[tokio::test]
1465    async fn test_empty_response() {
1466        let config = LlmSimConfig {
1467            response: ResponseConfig::Empty,
1468            tool_calls: None,
1469            simulate_latency: false,
1470            model_name: "test".to_string(),
1471            response_delay: None,
1472            response_id: None,
1473            effort_capture: None,
1474        };
1475
1476        let driver = LlmSimDriver::new(config);
1477        let messages = vec![user_message("test")];
1478
1479        let response = driver
1480            .chat_completion(messages, &make_config())
1481            .await
1482            .unwrap();
1483
1484        assert!(response.text.is_empty());
1485    }
1486
1487    #[test]
1488    fn test_driver_debug() {
1489        let driver = LlmSimDriver::new(LlmSimConfig::fixed("test").with_latency());
1490        let debug = format!("{:?}", driver);
1491
1492        assert!(debug.contains("LlmSimDriver"));
1493        assert!(debug.contains("simulate_latency"));
1494    }
1495
1496    #[test]
1497    fn test_default_config() {
1498        let config = LlmSimConfig::default();
1499        assert!(matches!(config.response, ResponseConfig::Fixed(_)));
1500        assert!(config.tool_calls.is_none());
1501        assert!(!config.simulate_latency);
1502    }
1503
1504    #[test]
1505    fn test_config_builder() {
1506        let tool_call = ToolCall {
1507            id: "call_1".to_string(),
1508            name: "get_weather".to_string(),
1509            arguments: serde_json::json!({"city": "NYC"}),
1510        };
1511
1512        let config = LlmSimConfig::fixed("Result")
1513            .with_tool_calls(vec![tool_call.clone()])
1514            .with_latency()
1515            .with_model("gpt-4")
1516            .with_response_delay(std::time::Duration::from_secs(2));
1517
1518        assert!(config.tool_calls.is_some());
1519        assert!(config.simulate_latency);
1520        assert_eq!(config.model_name, "gpt-4");
1521        assert_eq!(
1522            config.response_delay,
1523            Some(std::time::Duration::from_secs(2))
1524        );
1525    }
1526
1527    #[test]
1528    fn test_parse_ttft_from_model_name() {
1529        use super::parse_ttft_from_model_name;
1530
1531        // Valid patterns
1532        assert_eq!(
1533            parse_ttft_from_model_name("llmsim-ttft-2000"),
1534            Some(std::time::Duration::from_millis(2000))
1535        );
1536        assert_eq!(
1537            parse_ttft_from_model_name("test-ttft-500-extra"),
1538            Some(std::time::Duration::from_millis(500))
1539        );
1540
1541        // No TTFT patterns
1542        assert_eq!(parse_ttft_from_model_name("llmsim-model"), None);
1543        assert_eq!(parse_ttft_from_model_name("llmsim-ttft-0"), None);
1544        assert_eq!(parse_ttft_from_model_name("llmsim-ttft-abc"), None);
1545    }
1546
1547    #[test]
1548    fn test_resolve_latency_profile_from_model_name() {
1549        let driver = LlmSimDriver::new(LlmSimConfig::fixed("test"));
1550
1551        // "-latency" in model name -> fast profile (non-instant)
1552        let profile = driver.resolve_latency_profile("llmsim-latency");
1553        assert!(profile.sample_ttft().as_nanos() > 0);
1554
1555        // default model name -> instant profile
1556        let profile = driver.resolve_latency_profile("llmsim-default");
1557        assert_eq!(profile.sample_ttft().as_nanos(), 0);
1558
1559        // config flag also enables fast profile
1560        let driver = LlmSimDriver::new(LlmSimConfig::fixed("test").with_latency());
1561        let profile = driver.resolve_latency_profile("llmsim-default");
1562        assert!(profile.sample_ttft().as_nanos() > 0);
1563    }
1564
1565    #[tokio::test]
1566    async fn test_latency_streaming_from_model_name() {
1567        // Default driver (simulate_latency=false) but model name triggers latency
1568        let driver = LlmSimDriver::new(LlmSimConfig::fixed("Hello world"));
1569        let messages = vec![user_message("test")];
1570
1571        let mut config = make_config();
1572        config.model = "llmsim-latency".to_string();
1573
1574        let start = std::time::Instant::now();
1575        let mut stream = driver
1576            .chat_completion_stream(messages, &config)
1577            .await
1578            .unwrap();
1579
1580        let mut text_parts = Vec::new();
1581        let mut got_done = false;
1582
1583        while let Some(event) = stream.next().await {
1584            match event.unwrap() {
1585                LlmStreamEvent::TextDelta(text) => text_parts.push(text),
1586                LlmStreamEvent::Done(meta) => {
1587                    got_done = true;
1588                    assert_eq!(meta.model, Some("llmsim-latency".to_string()));
1589                }
1590                _ => {}
1591            }
1592        }
1593
1594        assert!(got_done);
1595        assert_eq!(text_parts.join(""), "Hello world");
1596        // With latency simulation, streaming should take non-zero time
1597        // (TTFT + inter-token delays)
1598        assert!(
1599            start.elapsed().as_millis() > 0,
1600            "latency simulation should introduce delays"
1601        );
1602    }
1603
1604    #[tokio::test]
1605    async fn test_no_latency_streaming_is_instant() {
1606        let driver = LlmSimDriver::new(LlmSimConfig::fixed("Hello world"));
1607        let messages = vec![user_message("test")];
1608
1609        let mut config = make_config();
1610        config.model = "llmsim-default".to_string();
1611
1612        let start = std::time::Instant::now();
1613        let response = driver.chat_completion(messages, &config).await.unwrap();
1614        let elapsed = start.elapsed();
1615
1616        assert_eq!(response.text, "Hello world");
1617        // Without latency, should complete nearly instantly (under 50ms)
1618        assert!(
1619            elapsed.as_millis() < 50,
1620            "instant mode should have no delays, took {}ms",
1621            elapsed.as_millis()
1622        );
1623    }
1624}
everruns_core/llmsim_driver.rs

everruns_core/
llmsim_driver.rs
