MockForge Architecture: How AI Features Integrate with a Rust-Based Mocking Framework

On this page

Building AI Features into a Rust Framework
The Core Architecture: Protocol Abstraction
Data Generation Pipeline
LLM Integration: Async and Non-Blocking
Caching Layer: Performance Optimization
Sync Service: Background Processing
Response Comparison: Deep Diff Algorithm
State Management: Concurrent Safe
Error Handling: Graceful Degradation
Performance Considerations
Observability: Built-In Instrumentation
The Bottom Line

🔗Building AI Features into a Rust Framework

MockForge is a Rust-based API mocking framework that supports HTTP, gRPC, GraphQL, WebSockets, and TCP. The core framework is built for performance and reliability—Rust's strengths. But modern mocking needs more than just protocol support; it needs intelligence.

This post covers the architectural decisions behind integrating AI features (LLM-powered data generation, automatic API sync) into a performance-focused Rust framework. How do you add AI capabilities without sacrificing the speed and reliability that make Rust attractive?

🔗The Core Architecture: Protocol Abstraction

MockForge's architecture is built around protocol abstraction. The core framework defines a unified interface for mock operations, and protocol-specific implementations (HTTP, gRPC, etc.) plug into this interface.

graph TD
    A[MockForge Core] -->|Protocol Abstraction| B[HTTP Handler]
    A -->|Protocol Abstraction| C[gRPC Handler]
    A -->|Protocol Abstraction| D[GraphQL Handler]
    A -->|Protocol Abstraction| E[WebSocket Handler]
    
    F[Data Generation] -->|Schema Input| A
    G[Sync Service] -->|Change Detection| A
    H[Response Selection] -->|Example Selection| A
    
    I[LLM Provider] -->|Generate Data| F
    J[Upstream API] -->|Poll & Compare| G

The AI features integrate at the data generation layer, not at the protocol layer. This means HTTP, gRPC, and GraphQL all benefit from LLM-powered generation without protocol-specific code.

🔗Data Generation Pipeline

The data generation pipeline is where AI features live. It's a pluggable system that supports multiple generation strategies:

pub enum GenerationStrategy {
    Example,      // Use schema examples
    Schema,       // Generate from schema structure
    LLM(LLMConfig), // Generate via LLM
    Faker,        // Use faker library
    Hybrid,       // Combine strategies
}

pub struct DataGenerator {
    strategy: GenerationStrategy,
    cache: Arc<Cache>,
    validator: SchemaValidator,
}

impl DataGenerator {
    pub async fn generate(&self, schema: &Value, context: &Context) -> Result<Value> {
        // Try strategy with caching
        match &self.strategy {
            GenerationStrategy::LLM(config) => {
                self.generate_via_llm(schema, context, config).await
            }
            GenerationStrategy::Schema => {
                self.generate_from_schema(schema).await
            }
            // ... other strategies
        }
    }
}

The generator is strategy-agnostic. It doesn't care whether data comes from an LLM or a faker—it just generates data that matches the schema. This makes it easy to swap strategies, fall back on errors, or combine strategies.

🔗LLM Integration: Async and Non-Blocking

Rust's async model makes LLM integration straightforward. LLM API calls are I/O-bound, so they fit naturally into async/await patterns.

pub struct LLMGenerator {
    client: reqwest::Client,
    config: LLMConfig,
    cache: Arc<Cache>,
}

impl LLMGenerator {
    pub async fn generate(&self, schema: &Value, context: &Context) -> Result<Value> {
        // Check cache first
        let cache_key = self.cache_key(schema, context);
        if let Some(cached) = self.cache.get(&cache_key).await? {
            return Ok(cached);
        }
        
        // Build prompt
        let prompt = self.build_prompt(schema, context)?;
        
        // Call LLM (non-blocking)
        let response = self.client
            .post(&self.config.api_url)
            .json(&prompt)
            .send()
            .await?;
        
        let generated: Value = response.json().await?;
        
        // Validate against schema
        self.validator.validate(&generated, schema)?;
        
        // Cache result
        self.cache.set(&cache_key, &generated).await?;
        
        Ok(generated)
    }
}

The async model means the mock server can handle other requests while waiting for LLM responses. This is crucial for performance—you don't want one slow LLM call blocking all other mock requests.

🔗Caching Layer: Performance Optimization

The caching layer is where performance meets cost optimization. It's a multi-tier cache:

pub struct Cache {
    memory: Arc<DashMap<String, CachedValue>>,  // In-memory (fast)
    disk: Arc<DiskCache>,                        // Disk (persistent)
    ttl: Duration,
}

impl Cache {
    pub async fn get(&self, key: &str) -> Option<Value> {
        // Try memory first
        if let Some(cached) = self.memory.get(key) {
            if !cached.is_expired() {
                return Some(cached.value.clone());
            }
        }
        
        // Try disk
        if let Some(cached) = self.disk.get(key).await? {
            if !cached.is_expired() {
                // Promote to memory
                self.memory.insert(key.to_string(), cached.clone());
                return Some(cached.value);
            }
        }
        
        None
    }
}

The multi-tier approach gives you:

Fast access: Memory cache for hot data
Persistence: Disk cache survives restarts
Cost savings: Avoid regenerating on every request

Cache keys are based on schema structure + context (seed, version, etc.), so identical schemas reuse the same cached data.

🔗Sync Service: Background Processing

The sync service runs as a background task, separate from the request handling path. This ensures sync operations don't impact mock server performance.

pub struct SyncService {
    config: Arc<RwLock<SyncConfig>>,
    database: Arc<RecorderDatabase>,
    http_client: Client,
    status: Arc<RwLock<SyncStatus>>,
}

impl SyncService {
    pub fn start(&self) -> JoinHandle<()> {
        tokio::spawn(async move {
            let mut interval = interval(Duration::from_secs(
                self.config.read().await.interval_seconds
            ));
            
            loop {
                interval.tick().await;
                
                // Sync doesn't block request handling
                let _ = self.sync_once().await;
            }
        })
    }
}

The sync service runs in its own tokio task, completely isolated from the HTTP server. This means:

Sync operations don't slow down mock responses
Sync failures don't crash the mock server
Sync can be paused/resumed without affecting mocks

🔗Response Comparison: Deep Diff Algorithm

The comparison algorithm is where the intelligence lives. It's not just string comparison—it's a content-aware diff engine.

pub struct ResponseComparator;

impl ResponseComparator {
    pub fn compare(
        original: &RecordedResponse,
        current: &RecordedResponse,
    ) -> ComparisonResult {
        let mut differences = Vec::new();
        
        // Compare status codes
        if original.status != current.status {
            differences.push(Difference::status_changed(...));
        }
        
        // Compare headers (ignoring dynamic ones)
        differences.extend(self.compare_headers(&original.headers, &current.headers));
        
        // Compare bodies (content-aware)
        match self.detect_content_type(&original.body) {
            ContentType::Json => {
                differences.extend(self.compare_json(&original.body, &current.body));
            }
            ContentType::Text => {
                differences.extend(self.compare_text(&original.body, &current.body));
            }
            _ => {}
        }
        
        ComparisonResult { matches: differences.is_empty(), differences }
    }
}

The comparison is smart about what to ignore (dynamic headers) and how to compare (JSON vs text). This gives you precise information about what changed without false positives from irrelevant differences.

🔗State Management: Concurrent Safe

Rust's ownership system makes concurrent state management safe. The sync service, cache, and mock handlers all share state, but Rust's type system prevents data races.

pub struct MockServer {
    // Shared, concurrent-safe state
    cache: Arc<Cache>,                    // Arc = shared ownership
    sync_service: Arc<SyncService>,       // Arc = shared ownership
    config: Arc<RwLock<Config>>,         // RwLock = read/write locking
    
    // Protocol handlers
    http_handler: HttpHandler,
    grpc_handler: GrpcHandler,
}

impl MockServer {
    pub async fn handle_request(&self, req: Request) -> Response {
        // All handlers can safely access shared state
        let data = self.cache.get(&req.schema_hash).await;
        let sync_status = self.sync_service.get_status().await;
        // ...
    }
}

Arc (Atomically Reference Counted) provides shared ownership without copying. RwLock provides read/write locking for concurrent access. Rust's type system ensures these are used correctly—you can't accidentally create a data race.

🔗Error Handling: Graceful Degradation

AI features can fail. The architecture handles this gracefully with fallback strategies.

pub async fn generate_with_fallbacks(
    generator: &DataGenerator,
    schema: &Value,
) -> Result<Value> {
    // Try LLM generation
    match generator.generate_via_llm(schema).await {
        Ok(data) => return Ok(data),
        Err(e) => {
            warn!("LLM generation failed: {}, falling back", e);
        }
    }
    
    // Try cached LLM data
    if let Some(cached) = generator.cache.get(schema).await {
        return Ok(cached);
    }
    
    // Fall back to schema-based generation
    generator.generate_from_schema(schema).await
}

The fallback chain ensures mocks always work, even when AI features fail. This is crucial for reliability—your tests shouldn't fail just because the LLM API is down.

🔗Performance Considerations

Adding AI features to a Rust framework doesn't mean sacrificing performance. The key is keeping AI operations out of the hot path.

Hot path (request handling):

Schema parsing: Fast (in-memory)
Cache lookup: Fast (in-memory hash map)
Response generation: Fast (cached or simple)

Cold path (background tasks):

LLM generation: Slow (API call, but async)
Sync operations: Slow (upstream API calls, but background)
Cache warming: Slow (but preemptive)

By keeping AI operations in the cold path, the mock server maintains sub-millisecond response times for cached requests, even with AI features enabled.

🔗Observability: Built-In Instrumentation

The architecture includes observability hooks at every layer:

#[instrument]
pub async fn generate_via_llm(schema: &Value) -> Result<Value> {
    let start = Instant::now();
    let result = self.call_llm_api(schema).await;
    let duration = start.elapsed();
    
    // Metrics
    metrics::llm_generation_duration.observe(duration.as_secs_f64());
    metrics::llm_generation_requests.inc();
    
    match &result {
        Ok(_) => metrics::llm_generation_success.inc(),
        Err(e) => {
            metrics::llm_generation_failure.inc();
            error!("LLM generation failed: {}", e);
        }
    }
    
    result
}

Tracing and metrics are built into the architecture, not bolted on. This makes it easy to monitor AI feature performance, costs, and reliability.

🔗The Bottom Line

Integrating AI features into a Rust framework is about architecture, not just implementation. The key decisions:

Separation of concerns: AI features in data generation layer, not protocol layer
Async model: Non-blocking LLM calls keep performance high
Caching strategy: Multi-tier cache balances performance and cost
Background processing: Sync operations don't impact request handling
Graceful degradation: Fallbacks ensure reliability when AI fails
Observability: Built-in instrumentation for monitoring and debugging

The result is a framework that has the performance of Rust with the intelligence of AI, without sacrificing either. That's the architectural challenge, and that's how MockForge solves it.