mockforge_ui/handlers/
behavioral_cloning.rs

1//! Behavioral cloning handlers for Admin UI
2//!
3//! This module provides API endpoints for managing flows and scenarios
4//! in the Admin UI.
5
6use axum::{
7    extract::{Path, Query, State},
8    response::Json,
9};
10use mockforge_recorder::behavioral_cloning::{
11    flow_recorder::{FlowRecorder, FlowRecordingConfig},
12    FlowCompiler, ScenarioStorage,
13};
14use mockforge_recorder::RecorderDatabase;
15use serde::Deserialize;
16use serde_json::{json, Value};
17use std::collections::HashMap;
18
19use crate::handlers::AdminState;
20use crate::models::ApiResponse;
21
22/// Get list of flows
23pub async fn get_flows(
24    State(_state): State<AdminState>,
25    Query(params): Query<HashMap<String, String>>,
26) -> Json<ApiResponse<Value>> {
27    // Get database path from config or use default
28    let db_path = params
29        .get("db_path")
30        .cloned()
31        .unwrap_or_else(|| "./mockforge-recordings.db".to_string());
32
33    let limit = params.get("limit").and_then(|s| s.parse::<usize>().ok()).unwrap_or(50);
34
35    match RecorderDatabase::new(&db_path).await {
36        Ok(db) => {
37            let recorder = FlowRecorder::new(db.clone(), FlowRecordingConfig::default());
38            match recorder.list_flows(Some(limit)).await {
39                Ok(flows) => {
40                    let flows_json: Vec<Value> = flows
41                        .into_iter()
42                        .map(|flow| {
43                            json!({
44                                "id": flow.id,
45                                "name": flow.name,
46                                "description": flow.description,
47                                "created_at": flow.created_at,
48                                "tags": flow.tags,
49                                "step_count": flow.steps.len(),
50                            })
51                        })
52                        .collect();
53
54                    Json(ApiResponse {
55                        success: true,
56                        data: Some(json!({
57                            "flows": flows_json,
58                            "total": flows_json.len()
59                        })),
60                        error: None,
61                        timestamp: chrono::Utc::now(),
62                    })
63                }
64                Err(e) => Json(ApiResponse {
65                    success: false,
66                    data: None,
67                    error: Some(format!("Failed to list flows: {}", e)),
68                    timestamp: chrono::Utc::now(),
69                }),
70            }
71        }
72        Err(e) => Json(ApiResponse {
73            success: false,
74            data: None,
75            error: Some(format!("Failed to connect to database: {}", e)),
76            timestamp: chrono::Utc::now(),
77        }),
78    }
79}
80
81/// Get flow details with timeline
82pub async fn get_flow(
83    State(_state): State<AdminState>,
84    Path(flow_id): Path<String>,
85    Query(params): Query<HashMap<String, String>>,
86) -> Json<ApiResponse<Value>> {
87    let db_path = params
88        .get("db_path")
89        .cloned()
90        .unwrap_or_else(|| "./mockforge-recordings.db".to_string());
91
92    match RecorderDatabase::new(&db_path).await {
93        Ok(db) => {
94            let recorder = FlowRecorder::new(db.clone(), FlowRecordingConfig::default());
95            match recorder.get_flow(&flow_id).await {
96                Ok(Some(flow)) => {
97                    // Build timeline data
98                    let steps: Vec<Value> = flow
99                        .steps
100                        .iter()
101                        .enumerate()
102                        .map(|(idx, step)| {
103                            json!({
104                                "index": idx,
105                                "request_id": step.request_id,
106                                "step_label": step.step_label,
107                                "timing_ms": step.timing_ms,
108                            })
109                        })
110                        .collect();
111
112                    Json(ApiResponse {
113                        success: true,
114                        data: Some(json!({
115                            "id": flow.id,
116                            "name": flow.name,
117                            "description": flow.description,
118                            "created_at": flow.created_at,
119                            "tags": flow.tags,
120                            "steps": steps,
121                            "step_count": steps.len(),
122                        })),
123                        error: None,
124                        timestamp: chrono::Utc::now(),
125                    })
126                }
127                Ok(None) => Json(ApiResponse {
128                    success: false,
129                    data: None,
130                    error: Some(format!("Flow not found: {}", flow_id)),
131                    timestamp: chrono::Utc::now(),
132                }),
133                Err(e) => Json(ApiResponse {
134                    success: false,
135                    data: None,
136                    error: Some(format!("Failed to get flow: {}", e)),
137                    timestamp: chrono::Utc::now(),
138                }),
139            }
140        }
141        Err(e) => Json(ApiResponse {
142            success: false,
143            data: None,
144            error: Some(format!("Failed to connect to database: {}", e)),
145            timestamp: chrono::Utc::now(),
146        }),
147    }
148}
149
150/// Tag a flow
151#[derive(Deserialize)]
152pub struct TagFlowRequest {
153    pub name: Option<String>,
154    pub description: Option<String>,
155    pub tags: Option<Vec<String>>,
156}
157
158pub async fn tag_flow(
159    State(_state): State<AdminState>,
160    Path(flow_id): Path<String>,
161    Query(params): Query<HashMap<String, String>>,
162    Json(payload): Json<TagFlowRequest>,
163) -> Json<ApiResponse<Value>> {
164    let db_path = params
165        .get("db_path")
166        .cloned()
167        .unwrap_or_else(|| "./mockforge-recordings.db".to_string());
168
169    match RecorderDatabase::new(&db_path).await {
170        Ok(db) => {
171            let recorder = FlowRecorder::new(db.clone(), FlowRecordingConfig::default());
172            match db
173                .update_flow_metadata(
174                    &flow_id,
175                    payload.name.as_deref(),
176                    payload.description.as_deref(),
177                    Some(&payload.tags.unwrap_or_default()),
178                )
179                .await
180            {
181                Ok(_) => Json(ApiResponse {
182                    success: true,
183                    data: Some(json!({
184                        "message": "Flow tagged successfully",
185                        "flow_id": flow_id
186                    })),
187                    error: None,
188                    timestamp: chrono::Utc::now(),
189                }),
190                Err(e) => Json(ApiResponse {
191                    success: false,
192                    data: None,
193                    error: Some(format!("Failed to tag flow: {}", e)),
194                    timestamp: chrono::Utc::now(),
195                }),
196            }
197        }
198        Err(e) => Json(ApiResponse {
199            success: false,
200            data: None,
201            error: Some(format!("Failed to connect to database: {}", e)),
202            timestamp: chrono::Utc::now(),
203        }),
204    }
205}
206
207/// Compile flow to scenario
208#[derive(Deserialize)]
209pub struct CompileFlowRequest {
210    pub scenario_name: String,
211    pub flex_mode: Option<bool>,
212}
213
214pub async fn compile_flow(
215    State(_state): State<AdminState>,
216    Path(flow_id): Path<String>,
217    Query(params): Query<HashMap<String, String>>,
218    Json(payload): Json<CompileFlowRequest>,
219) -> Json<ApiResponse<Value>> {
220    let db_path = params
221        .get("db_path")
222        .cloned()
223        .unwrap_or_else(|| "./mockforge-recordings.db".to_string());
224
225    match RecorderDatabase::new(&db_path).await {
226        Ok(db) => {
227            let recorder = FlowRecorder::new(db.clone(), FlowRecordingConfig::default());
228            match recorder.get_flow(&flow_id).await {
229                Ok(Some(flow)) => {
230                    let compiler = FlowCompiler::new(db.clone());
231                    let strict_mode = !payload.flex_mode.unwrap_or(false);
232                    match compiler
233                        .compile_flow(&flow, payload.scenario_name.clone(), strict_mode)
234                        .await
235                    {
236                        Ok(scenario) => {
237                            // Store the scenario
238                            let storage = ScenarioStorage::new(db);
239                            match storage.store_scenario_auto_version(&scenario).await {
240                                Ok(version) => Json(ApiResponse {
241                                    success: true,
242                                    data: Some(json!({
243                                        "scenario_id": scenario.id,
244                                        "scenario_name": scenario.name,
245                                        "version": version,
246                                        "message": "Flow compiled successfully"
247                                    })),
248                                    error: None,
249                                    timestamp: chrono::Utc::now(),
250                                }),
251                                Err(e) => Json(ApiResponse {
252                                    success: false,
253                                    data: None,
254                                    error: Some(format!("Failed to store scenario: {}", e)),
255                                    timestamp: chrono::Utc::now(),
256                                }),
257                            }
258                        }
259                        Err(e) => Json(ApiResponse {
260                            success: false,
261                            data: None,
262                            error: Some(format!("Failed to compile flow: {}", e)),
263                            timestamp: chrono::Utc::now(),
264                        }),
265                    }
266                }
267                Ok(None) => Json(ApiResponse {
268                    success: false,
269                    data: None,
270                    error: Some(format!("Flow not found: {}", flow_id)),
271                    timestamp: chrono::Utc::now(),
272                }),
273                Err(e) => Json(ApiResponse {
274                    success: false,
275                    data: None,
276                    error: Some(format!("Failed to get flow: {}", e)),
277                    timestamp: chrono::Utc::now(),
278                }),
279            }
280        }
281        Err(e) => Json(ApiResponse {
282            success: false,
283            data: None,
284            error: Some(format!("Failed to connect to database: {}", e)),
285            timestamp: chrono::Utc::now(),
286        }),
287    }
288}
289
290/// Get list of scenarios
291pub async fn get_scenarios(
292    State(_state): State<AdminState>,
293    Query(params): Query<HashMap<String, String>>,
294) -> Json<ApiResponse<Value>> {
295    let db_path = params
296        .get("db_path")
297        .cloned()
298        .unwrap_or_else(|| "./mockforge-recordings.db".to_string());
299
300    let limit = params.get("limit").and_then(|s| s.parse::<usize>().ok()).unwrap_or(50);
301
302    match RecorderDatabase::new(&db_path).await {
303        Ok(db) => {
304            let storage = ScenarioStorage::new(db);
305            match storage.list_scenarios(Some(limit)).await {
306                Ok(scenarios) => {
307                    let scenarios_json: Vec<Value> = scenarios
308                        .into_iter()
309                        .map(|s| {
310                            json!({
311                                "id": s.id,
312                                "name": s.name,
313                                "version": s.version,
314                                "description": s.description,
315                                "created_at": s.created_at,
316                                "updated_at": s.updated_at,
317                                "tags": s.tags,
318                            })
319                        })
320                        .collect();
321
322                    Json(ApiResponse {
323                        success: true,
324                        data: Some(json!({
325                            "scenarios": scenarios_json,
326                            "total": scenarios_json.len()
327                        })),
328                        error: None,
329                        timestamp: chrono::Utc::now(),
330                    })
331                }
332                Err(e) => Json(ApiResponse {
333                    success: false,
334                    data: None,
335                    error: Some(format!("Failed to list scenarios: {}", e)),
336                    timestamp: chrono::Utc::now(),
337                }),
338            }
339        }
340        Err(e) => Json(ApiResponse {
341            success: false,
342            data: None,
343            error: Some(format!("Failed to connect to database: {}", e)),
344            timestamp: chrono::Utc::now(),
345        }),
346    }
347}
348
349/// Get scenario details
350pub async fn get_scenario(
351    State(_state): State<AdminState>,
352    Path(scenario_id): Path<String>,
353    Query(params): Query<HashMap<String, String>>,
354) -> Json<ApiResponse<Value>> {
355    let db_path = params
356        .get("db_path")
357        .cloned()
358        .unwrap_or_else(|| "./mockforge-recordings.db".to_string());
359
360    match RecorderDatabase::new(&db_path).await {
361        Ok(db) => {
362            let storage = ScenarioStorage::new(db);
363            match storage.get_scenario(&scenario_id).await {
364                Ok(Some(scenario)) => {
365                    let steps: Vec<Value> = scenario
366                        .steps
367                        .iter()
368                        .map(|step| {
369                            json!({
370                                "step_id": step.step_id,
371                                "label": step.label,
372                                "method": step.request.method,
373                                "path": step.request.path,
374                                "status_code": step.response.status_code,
375                                "timing_ms": step.timing_ms,
376                            })
377                        })
378                        .collect();
379
380                    Json(ApiResponse {
381                        success: true,
382                        data: Some(json!({
383                            "id": scenario.id,
384                            "name": scenario.name,
385                            "description": scenario.description,
386                            "strict_mode": scenario.strict_mode,
387                            "steps": steps,
388                            "step_count": steps.len(),
389                            "state_variables": scenario.state_variables.len(),
390                        })),
391                        error: None,
392                        timestamp: chrono::Utc::now(),
393                    })
394                }
395                Ok(None) => Json(ApiResponse {
396                    success: false,
397                    data: None,
398                    error: Some(format!("Scenario not found: {}", scenario_id)),
399                    timestamp: chrono::Utc::now(),
400                }),
401                Err(e) => Json(ApiResponse {
402                    success: false,
403                    data: None,
404                    error: Some(format!("Failed to get scenario: {}", e)),
405                    timestamp: chrono::Utc::now(),
406                }),
407            }
408        }
409        Err(e) => Json(ApiResponse {
410            success: false,
411            data: None,
412            error: Some(format!("Failed to connect to database: {}", e)),
413            timestamp: chrono::Utc::now(),
414        }),
415    }
416}
417
418/// Export scenario
419pub async fn export_scenario(
420    State(_state): State<AdminState>,
421    Path(scenario_id): Path<String>,
422    Query(params): Query<HashMap<String, String>>,
423) -> Json<ApiResponse<Value>> {
424    let db_path = params
425        .get("db_path")
426        .cloned()
427        .unwrap_or_else(|| "./mockforge-recordings.db".to_string());
428
429    let format = params.get("format").cloned().unwrap_or_else(|| "yaml".to_string());
430
431    match RecorderDatabase::new(&db_path).await {
432        Ok(db) => {
433            let storage = ScenarioStorage::new(db);
434            match storage.export_scenario(&scenario_id, &format).await {
435                Ok(content) => Json(ApiResponse {
436                    success: true,
437                    data: Some(json!({
438                        "scenario_id": scenario_id,
439                        "format": format,
440                        "content": content,
441                    })),
442                    error: None,
443                    timestamp: chrono::Utc::now(),
444                }),
445                Err(e) => Json(ApiResponse {
446                    success: false,
447                    data: None,
448                    error: Some(format!("Failed to export scenario: {}", e)),
449                    timestamp: chrono::Utc::now(),
450                }),
451            }
452        }
453        Err(e) => Json(ApiResponse {
454            success: false,
455            data: None,
456            error: Some(format!("Failed to connect to database: {}", e)),
457            timestamp: chrono::Utc::now(),
458        }),
459    }
460}