use super::config::IntelligentBehaviorConfig;
use super::context::StatefulAiContext;
use super::mutation_analyzer::MutationAnalyzer;
use super::pagination_intelligence::{
PaginationIntelligence, PaginationMetadata, PaginationRequest,
};
use super::rule_generator::{ExamplePair, RuleGenerator};
use super::types::BehaviorRules;
use super::validation_generator::{RequestContext, ValidationGenerator};
use crate::openapi::OpenApiSpec;
use crate::Result;
use serde_json::Value;
use std::collections::HashMap;
use uuid;
#[derive(Debug, Clone)]
pub struct Request {
pub method: String,
pub path: String,
pub body: Option<Value>,
pub query_params: HashMap<String, String>,
pub headers: HashMap<String, String>,
}
#[derive(Debug, Clone)]
pub struct Response {
pub status_code: u16,
pub body: Value,
pub headers: HashMap<String, String>,
}
pub struct MockAI {
rules: BehaviorRules,
rule_generator: RuleGenerator,
mutation_analyzer: MutationAnalyzer,
validation_generator: ValidationGenerator,
pagination_intelligence: PaginationIntelligence,
config: IntelligentBehaviorConfig,
session_contexts: std::sync::Arc<tokio::sync::RwLock<HashMap<String, StatefulAiContext>>>,
}
impl MockAI {
pub async fn from_openapi(
spec: &OpenApiSpec,
config: IntelligentBehaviorConfig,
) -> Result<Self> {
let examples = Self::extract_examples_from_openapi(spec)?;
let behavior_config = config.behavior_model.clone();
let rule_generator = RuleGenerator::new(behavior_config.clone());
let rules = rule_generator.generate_rules_from_examples(examples).await?;
let mutation_analyzer = MutationAnalyzer::new().with_rules(rules.clone());
let validation_generator = ValidationGenerator::new(behavior_config.clone());
let pagination_intelligence = PaginationIntelligence::new(behavior_config);
Ok(Self {
rules,
rule_generator,
mutation_analyzer,
validation_generator,
pagination_intelligence,
config,
session_contexts: std::sync::Arc::new(tokio::sync::RwLock::new(HashMap::new())),
})
}
pub async fn from_examples(
examples: Vec<ExamplePair>,
config: IntelligentBehaviorConfig,
) -> Result<Self> {
let behavior_config = config.behavior_model.clone();
let rule_generator = RuleGenerator::new(behavior_config.clone());
let rules = rule_generator.generate_rules_from_examples(examples).await?;
let mutation_analyzer = MutationAnalyzer::new().with_rules(rules.clone());
let validation_generator = ValidationGenerator::new(behavior_config.clone());
let pagination_intelligence = PaginationIntelligence::new(behavior_config);
Ok(Self {
rules,
rule_generator,
mutation_analyzer,
validation_generator,
pagination_intelligence,
config,
session_contexts: std::sync::Arc::new(tokio::sync::RwLock::new(HashMap::new())),
})
}
pub fn new(config: IntelligentBehaviorConfig) -> Self {
let behavior_config = config.behavior_model.clone();
let rule_generator = RuleGenerator::new(behavior_config.clone());
let rules = BehaviorRules::default();
let mutation_analyzer = MutationAnalyzer::new().with_rules(rules.clone());
let validation_generator = ValidationGenerator::new(behavior_config.clone());
let pagination_intelligence = PaginationIntelligence::new(behavior_config);
Self {
rules,
rule_generator,
mutation_analyzer,
validation_generator,
pagination_intelligence,
config,
session_contexts: std::sync::Arc::new(tokio::sync::RwLock::new(HashMap::new())),
}
}
pub async fn process_request(&self, request: &Request) -> Result<Response> {
let session_id = self.extract_session_id(request);
let session_context = self.get_or_create_session_context(session_id).await?;
let response = self.generate_response(request, &session_context).await?;
if let Err(e) = session_context
.record_interaction(
request.method.clone(),
request.path.clone(),
request.body.clone(),
Some(response.body.clone()),
)
.await
{
tracing::warn!("Failed to record interaction: {}", e);
}
Ok(response)
}
fn extract_session_id(&self, request: &Request) -> Option<String> {
if let Some(session_id) = request.headers.get("X-Session-ID") {
return Some(session_id.clone());
}
if let Some(cookie_header) = request.headers.get("Cookie") {
for part in cookie_header.split(';') {
let part = part.trim();
if let Some((key, value)) = part.split_once('=') {
if key.trim() == "mockforge_session" {
return Some(value.trim().to_string());
}
}
}
}
None
}
async fn get_or_create_session_context(
&self,
session_id: Option<String>,
) -> Result<StatefulAiContext> {
let session_id = session_id.unwrap_or_else(|| format!("session_{}", uuid::Uuid::new_v4()));
{
let contexts = self.session_contexts.read().await;
if let Some(context) = contexts.get(&session_id) {
return Ok(context.clone());
}
}
let new_context = StatefulAiContext::new(session_id.clone(), self.config.clone());
{
let mut contexts = self.session_contexts.write().await;
contexts.insert(session_id, new_context.clone());
}
Ok(new_context)
}
pub async fn generate_response(
&self,
request: &Request,
session_context: &StatefulAiContext,
) -> Result<Response> {
let method_upper = request.method.to_uppercase();
let is_mutation_method =
matches!(method_upper.as_str(), "POST" | "PUT" | "PATCH" | "DELETE");
let history = session_context.get_history().await;
let previous_request = history.last().and_then(|interaction| interaction.request.clone());
let mutation_analysis = if is_mutation_method {
let current_body = request.body.clone().unwrap_or(serde_json::json!({}));
self.mutation_analyzer
.analyze_mutation(¤t_body, previous_request.as_ref(), session_context)
.await?
} else {
super::mutation_analyzer::MutationAnalysis {
mutation_type: super::mutation_analyzer::MutationType::NoChange, changed_fields: Vec::new(),
added_fields: Vec::new(),
removed_fields: Vec::new(),
validation_issues: Vec::new(),
confidence: 1.0,
}
};
if !mutation_analysis.validation_issues.is_empty() {
let issue = &mutation_analysis.validation_issues[0];
let request_context = RequestContext {
method: request.method.clone(),
path: request.path.clone(),
request_body: request.body.clone(),
query_params: request.query_params.clone(),
headers: request.headers.clone(),
};
let error_response = self
.validation_generator
.generate_validation_error(issue, &request_context)
.await?;
return Ok(Response {
status_code: error_response.status_code,
body: error_response.body,
headers: HashMap::new(),
});
}
if self.is_paginated_request(request) {
let pagination_meta =
self.generate_pagination_metadata(request, session_context).await?;
let body = self.build_paginated_response(&pagination_meta, request).await?;
return Ok(Response {
status_code: 200,
body,
headers: HashMap::new(),
});
}
let response_body = if is_mutation_method {
self.generate_response_body(&mutation_analysis, request, session_context)
.await?
} else {
tracing::debug!(
"Skipping mutation-based response generation for {} request - using OpenAPI response generation",
method_upper
);
serde_json::json!({}) };
Ok(Response {
status_code: 200,
body: response_body,
headers: HashMap::new(),
})
}
pub async fn learn_from_example(&mut self, example: ExamplePair) -> Result<()> {
let examples = vec![example];
let new_rules = self.rule_generator.generate_rules_from_examples(examples).await?;
self.merge_rules(new_rules);
Ok(())
}
pub fn rules(&self) -> &BehaviorRules {
&self.rules
}
pub fn update_rules(&mut self, rules: BehaviorRules) {
self.rules = rules;
self.mutation_analyzer = MutationAnalyzer::new().with_rules(self.rules.clone());
}
pub fn update_config(&mut self, config: IntelligentBehaviorConfig) {
self.config = config.clone();
let behavior_config = self.config.behavior_model.clone();
self.validation_generator = ValidationGenerator::new(behavior_config.clone());
self.pagination_intelligence = PaginationIntelligence::new(behavior_config);
}
pub async fn update_config_async(
this: &std::sync::Arc<tokio::sync::RwLock<Self>>,
config: IntelligentBehaviorConfig,
) -> Result<()> {
let mut mockai = this.write().await;
mockai.update_config(config);
Ok(())
}
pub fn get_config(&self) -> &IntelligentBehaviorConfig {
&self.config
}
pub fn extract_examples_from_openapi(spec: &OpenApiSpec) -> Result<Vec<ExamplePair>> {
let mut examples = Vec::new();
let path_operations = spec.all_paths_and_operations();
for (path, operations) in path_operations {
for (method, operation) in operations {
let request = operation
.request_body
.as_ref()
.and_then(|rb| rb.as_item())
.and_then(|rb| rb.content.get("application/json"))
.and_then(|media| media.example.clone());
let response = operation.responses.responses.iter().find_map(|(status, resp)| {
if let openapiv3::StatusCode::Code(200) = status {
resp.as_item()
.and_then(|r| r.content.get("application/json"))
.and_then(|media| media.example.clone())
} else {
None
}
});
examples.push(ExamplePair {
method: method.clone(),
path: path.clone(),
request,
status: 200,
response,
query_params: HashMap::new(),
headers: HashMap::new(),
metadata: HashMap::new(),
});
}
}
Ok(examples)
}
fn is_paginated_request(&self, request: &Request) -> bool {
request.query_params.keys().any(|key| {
matches!(
key.to_lowercase().as_str(),
"page" | "limit" | "per_page" | "offset" | "cursor"
)
})
}
async fn generate_pagination_metadata(
&self,
request: &Request,
session_context: &StatefulAiContext,
) -> Result<PaginationMetadata> {
let pagination_request = PaginationRequest {
path: request.path.clone(),
query_params: request.query_params.clone(),
request_body: request.body.clone(),
};
self.pagination_intelligence
.generate_pagination_metadata(&pagination_request, session_context)
.await
}
async fn build_paginated_response(
&self,
meta: &PaginationMetadata,
_request: &Request,
) -> Result<Value> {
Ok(serde_json::json!({
"data": [], "pagination": {
"page": meta.page,
"page_size": meta.page_size,
"total": meta.total,
"total_pages": meta.total_pages,
"has_next": meta.has_next,
"has_prev": meta.has_prev,
"offset": meta.offset,
"next_cursor": meta.next_cursor,
"prev_cursor": meta.prev_cursor,
}
}))
}
async fn generate_response_body(
&self,
mutation: &super::mutation_analyzer::MutationAnalysis,
request: &Request,
_session_context: &StatefulAiContext,
) -> Result<Value> {
match mutation.mutation_type {
super::mutation_analyzer::MutationType::NoChange => {
tracing::debug!("MutationType::NoChange - returning empty object to use OpenAPI response generation");
Ok(serde_json::json!({}))
}
super::mutation_analyzer::MutationType::Create => {
Ok(serde_json::json!({
"id": "generated_id",
"status": "created",
"data": request.body.clone().unwrap_or(serde_json::json!({}))
}))
}
super::mutation_analyzer::MutationType::Update
| super::mutation_analyzer::MutationType::PartialUpdate => {
Ok(serde_json::json!({
"id": "resource_id",
"status": "updated",
"data": request.body.clone().unwrap_or(serde_json::json!({}))
}))
}
super::mutation_analyzer::MutationType::Delete => {
Ok(serde_json::json!({
"status": "deleted",
"message": "Resource deleted successfully"
}))
}
_ => {
Ok(serde_json::json!({
"status": "success",
"data": request.body.clone().unwrap_or(serde_json::json!({}))
}))
}
}
}
fn merge_rules(&mut self, new_rules: BehaviorRules) {
self.rules.consistency_rules.extend(new_rules.consistency_rules);
for (key, value) in new_rules.schemas {
self.rules.schemas.insert(key, value);
}
for (key, value) in new_rules.state_transitions {
self.rules.state_transitions.insert(key, value);
}
if new_rules.system_prompt.len() > self.rules.system_prompt.len() {
self.rules.system_prompt = new_rules.system_prompt;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[tokio::test]
async fn test_is_paginated_request() {
if std::env::var("OPENAI_API_KEY").is_err() && std::env::var("ANTHROPIC_API_KEY").is_err() {
eprintln!("Skipping test: No API key found");
return;
}
let config = IntelligentBehaviorConfig::default();
let examples = vec![ExamplePair {
method: "GET".to_string(),
path: "/api/users".to_string(),
request: None,
status: 200,
response: Some(json!({})),
query_params: HashMap::new(),
headers: HashMap::new(),
metadata: HashMap::new(),
}];
let mockai = match MockAI::from_examples(examples, config).await {
Ok(m) => m,
Err(e) => {
eprintln!("Skipping test: Failed to create MockAI: {}", e);
return;
}
};
let mut query_params = HashMap::new();
query_params.insert("page".to_string(), "1".to_string());
let request = Request {
method: "GET".to_string(),
path: "/api/users".to_string(),
body: None,
query_params,
headers: HashMap::new(),
};
assert!(mockai.is_paginated_request(&request));
}
#[tokio::test]
async fn test_process_request() {
if std::env::var("OPENAI_API_KEY").is_err() && std::env::var("ANTHROPIC_API_KEY").is_err() {
eprintln!("Skipping test: No API key found");
return;
}
let config = IntelligentBehaviorConfig::default();
let examples = vec![ExamplePair {
method: "GET".to_string(),
path: "/api/users".to_string(),
request: None,
status: 200,
response: Some(json!({
"users": [],
"total": 0
})),
query_params: HashMap::new(),
headers: HashMap::new(),
metadata: HashMap::new(),
}];
let mockai = match MockAI::from_examples(examples, config).await {
Ok(m) => m,
Err(e) => {
eprintln!("Skipping test: Failed to create MockAI: {}", e);
return;
}
};
let request = Request {
method: "GET".to_string(),
path: "/api/users".to_string(),
body: None,
query_params: HashMap::new(),
headers: HashMap::new(),
};
let response = match mockai.process_request(&request).await {
Ok(r) => r,
Err(e) => {
eprintln!("Skipping test: Failed to process request: {}", e);
return;
}
};
assert_eq!(response.status_code, 200);
assert!(response.body.is_object() || response.body.is_array());
}
#[tokio::test]
async fn test_process_request_with_body() {
if std::env::var("OPENAI_API_KEY").is_err() && std::env::var("ANTHROPIC_API_KEY").is_err() {
eprintln!("Skipping test: No API key found");
return;
}
let config = IntelligentBehaviorConfig::default();
let examples = vec![ExamplePair {
method: "POST".to_string(),
path: "/api/users".to_string(),
request: Some(json!({
"name": "John Doe",
"email": "john@example.com"
})),
status: 201,
response: Some(json!({
"id": "123",
"name": "John Doe",
"email": "john@example.com"
})),
query_params: HashMap::new(),
headers: HashMap::new(),
metadata: HashMap::new(),
}];
let mockai = match MockAI::from_examples(examples, config).await {
Ok(m) => m,
Err(e) => {
eprintln!("Skipping test: Failed to create MockAI: {}", e);
return;
}
};
let request = Request {
method: "POST".to_string(),
path: "/api/users".to_string(),
body: Some(json!({
"name": "Jane Doe",
"email": "jane@example.com"
})),
query_params: HashMap::new(),
headers: HashMap::new(),
};
let response = match mockai.process_request(&request).await {
Ok(r) => r,
Err(e) => {
eprintln!("Skipping test: Failed to process request: {}", e);
return;
}
};
assert_eq!(response.status_code, 201);
assert!(response.body.is_object());
}
#[test]
fn test_request_creation() {
let mut query_params = HashMap::new();
query_params.insert("page".to_string(), "1".to_string());
let mut headers = HashMap::new();
headers.insert("Content-Type".to_string(), "application/json".to_string());
let request = Request {
method: "GET".to_string(),
path: "/api/users".to_string(),
body: Some(json!({"id": 1})),
query_params,
headers,
};
assert_eq!(request.method, "GET");
assert_eq!(request.path, "/api/users");
assert!(request.body.is_some());
}
#[test]
fn test_response_creation() {
let mut headers = HashMap::new();
headers.insert("Content-Type".to_string(), "application/json".to_string());
let response = Response {
status_code: 200,
body: json!({"message": "success"}),
headers,
};
assert_eq!(response.status_code, 200);
assert!(response.body.is_object());
}
#[test]
fn test_mockai_new() {
let config = IntelligentBehaviorConfig::default();
let mockai = MockAI::new(config);
let _ = mockai;
}
#[test]
fn test_mockai_rules() {
let config = IntelligentBehaviorConfig::default();
let mockai = MockAI::new(config);
let rules = mockai.rules();
let _ = rules;
}
#[test]
fn test_mockai_update_rules() {
let config = IntelligentBehaviorConfig::default();
let mut mockai = MockAI::new(config);
let new_rules = BehaviorRules::default();
mockai.update_rules(new_rules);
}
#[test]
fn test_mockai_get_config() {
let config = IntelligentBehaviorConfig::default();
let mockai = MockAI::new(config.clone());
let retrieved_config = mockai.get_config();
let _ = retrieved_config;
}
#[test]
fn test_mockai_update_config() {
let config = IntelligentBehaviorConfig::default();
let mut mockai = MockAI::new(config.clone());
let new_config = IntelligentBehaviorConfig::default();
mockai.update_config(new_config);
}
#[test]
fn test_extract_examples_from_openapi_empty_spec() {
let spec_json = json!({
"openapi": "3.0.0",
"info": {
"title": "Test API",
"version": "1.0.0"
},
"paths": {}
});
let spec = OpenApiSpec::from_json(spec_json).unwrap();
let examples = MockAI::extract_examples_from_openapi(&spec).unwrap();
assert!(examples.is_empty());
}
#[test]
fn test_request_with_all_fields() {
let mut headers = HashMap::new();
headers.insert("Authorization".to_string(), "Bearer token".to_string());
let mut query_params = HashMap::new();
query_params.insert("limit".to_string(), "10".to_string());
let request = Request {
method: "POST".to_string(),
path: "/api/data".to_string(),
body: Some(json!({"key": "value"})),
query_params: query_params.clone(),
headers: headers.clone(),
};
assert_eq!(request.method, "POST");
assert_eq!(request.path, "/api/data");
assert!(request.body.is_some());
assert_eq!(request.query_params.get("limit"), Some(&"10".to_string()));
assert_eq!(request.headers.get("Authorization"), Some(&"Bearer token".to_string()));
}
#[test]
fn test_response_with_headers() {
let mut headers = HashMap::new();
headers.insert("X-Total-Count".to_string(), "100".to_string());
headers.insert("Content-Type".to_string(), "application/json".to_string());
let response = Response {
status_code: 201,
body: json!({"id": "123", "created": true}),
headers: headers.clone(),
};
assert_eq!(response.status_code, 201);
assert!(response.body.is_object());
assert_eq!(response.headers.len(), 2);
assert_eq!(response.headers.get("X-Total-Count"), Some(&"100".to_string()));
}
#[test]
fn test_request_clone() {
let request1 = Request {
method: "GET".to_string(),
path: "/api/test".to_string(),
body: Some(json!({"id": 1})),
query_params: HashMap::new(),
headers: HashMap::new(),
};
let request2 = request1.clone();
assert_eq!(request1.method, request2.method);
assert_eq!(request1.path, request2.path);
}
#[test]
fn test_request_debug() {
let request = Request {
method: "POST".to_string(),
path: "/api/users".to_string(),
body: None,
query_params: HashMap::new(),
headers: HashMap::new(),
};
let debug_str = format!("{:?}", request);
assert!(debug_str.contains("Request"));
}
#[test]
fn test_response_clone() {
let response1 = Response {
status_code: 200,
body: json!({"status": "ok"}),
headers: HashMap::new(),
};
let response2 = response1.clone();
assert_eq!(response1.status_code, response2.status_code);
}
#[test]
fn test_response_debug() {
let response = Response {
status_code: 404,
body: json!({"error": "Not found"}),
headers: HashMap::new(),
};
let debug_str = format!("{:?}", response);
assert!(debug_str.contains("Response"));
}
#[test]
fn test_request_with_empty_fields() {
let request = Request {
method: "GET".to_string(),
path: "/api/test".to_string(),
body: None,
query_params: HashMap::new(),
headers: HashMap::new(),
};
assert!(request.body.is_none());
assert!(request.query_params.is_empty());
assert!(request.headers.is_empty());
}
#[test]
fn test_response_with_empty_headers() {
let response = Response {
status_code: 200,
body: json!({"data": []}),
headers: HashMap::new(),
};
assert!(response.headers.is_empty());
assert!(response.body.is_object());
}
#[test]
fn test_request_with_complex_body() {
let request = Request {
method: "PUT".to_string(),
path: "/api/users/123".to_string(),
body: Some(json!({
"name": "John Doe",
"email": "john@example.com",
"metadata": {
"role": "admin",
"permissions": ["read", "write"]
}
})),
query_params: HashMap::new(),
headers: HashMap::new(),
};
assert!(request.body.is_some());
let body = request.body.unwrap();
assert!(body.is_object());
assert!(body.get("metadata").is_some());
}
#[test]
fn test_response_with_array_body() {
let response = Response {
status_code: 200,
body: json!([
{"id": 1, "name": "Alice"},
{"id": 2, "name": "Bob"},
{"id": 3, "name": "Charlie"}
]),
headers: HashMap::new(),
};
assert!(response.body.is_array());
let array = response.body.as_array().unwrap();
assert_eq!(array.len(), 3);
}
#[test]
fn test_request_with_multiple_query_params() {
let mut query_params = HashMap::new();
query_params.insert("page".to_string(), "1".to_string());
query_params.insert("limit".to_string(), "20".to_string());
query_params.insert("sort".to_string(), "name".to_string());
query_params.insert("order".to_string(), "asc".to_string());
let request = Request {
method: "GET".to_string(),
path: "/api/users".to_string(),
body: None,
query_params: query_params.clone(),
headers: HashMap::new(),
};
assert_eq!(request.query_params.len(), 4);
assert_eq!(request.query_params.get("page"), Some(&"1".to_string()));
assert_eq!(request.query_params.get("limit"), Some(&"20".to_string()));
}
#[test]
fn test_response_with_multiple_headers() {
let mut headers = HashMap::new();
headers.insert("Content-Type".to_string(), "application/json".to_string());
headers.insert("X-Request-ID".to_string(), "req-123".to_string());
headers.insert("X-Rate-Limit-Remaining".to_string(), "99".to_string());
headers.insert("Cache-Control".to_string(), "no-cache".to_string());
let response = Response {
status_code: 200,
body: json!({"data": "test"}),
headers: headers.clone(),
};
assert_eq!(response.headers.len(), 4);
assert_eq!(response.headers.get("X-Request-ID"), Some(&"req-123".to_string()));
}
#[test]
fn test_request_different_methods() {
let methods = vec!["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"];
for method in methods {
let request = Request {
method: method.to_string(),
path: "/api/test".to_string(),
body: None,
query_params: HashMap::new(),
headers: HashMap::new(),
};
assert_eq!(request.method, method);
}
}
#[test]
fn test_response_different_status_codes() {
let status_codes = vec![200, 201, 204, 400, 401, 403, 404, 500, 503];
for status_code in status_codes {
let response = Response {
status_code,
body: json!({"status": status_code}),
headers: HashMap::new(),
};
assert_eq!(response.status_code, status_code);
}
}
}