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

everruns_core/
llmsim_driver.rs
