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

everruns_core/
llmsim_driver.rs
