1use axum::{
7 extract::{Path, Query, State},
8 http::StatusCode,
9 response::Json,
10};
11use mockforge_core::consistency::ConsistencyEngine;
12use serde::Deserialize;
13use serde_json::Value;
14use std::collections::HashMap;
15use std::sync::Arc;
16use tokio::sync::RwLock;
17use tracing::debug;
18
19#[derive(Debug, Clone)]
21pub(crate) struct RequestContextSnapshot {
22 workspace_id: String,
24 state_snapshot: serde_json::Value,
26 timestamp: i64,
28}
29
30#[derive(Clone)]
32pub struct XRayState {
33 pub engine: Arc<ConsistencyEngine>,
35 pub request_contexts: Arc<RwLock<HashMap<String, RequestContextSnapshot>>>,
37}
38
39#[derive(Debug, Deserialize)]
41pub struct XRayQuery {
42 #[serde(default = "default_workspace")]
44 pub workspace: String,
45}
46
47fn default_workspace() -> String {
48 "default".to_string()
49}
50
51pub async fn get_state_summary(
57 State(state): State<XRayState>,
58 Query(params): Query<XRayQuery>,
59) -> Result<Json<Value>, StatusCode> {
60 let unified_state = state.engine.get_state(¶ms.workspace).await.ok_or_else(|| {
61 debug!("No state found for workspace: {}", params.workspace);
62 StatusCode::NOT_FOUND
63 })?;
64
65 let summary = serde_json::json!({
67 "workspace_id": unified_state.workspace_id,
68 "scenario": unified_state.active_scenario,
69 "persona": unified_state.active_persona.as_ref().map(|p| serde_json::json!({
70 "id": p.id,
71 "traits": p.traits,
72 })),
73 "reality_level": unified_state.reality_level.value(),
74 "reality_level_name": unified_state.reality_level.name(),
75 "reality_ratio": unified_state.reality_continuum_ratio,
76 "chaos_rules": unified_state
78 .active_chaos_rules
79 .iter()
80 .filter_map(|r| r.get("name").and_then(|v| v.as_str()).map(|s| s.to_string()))
81 .collect::<Vec<_>>(),
82 "timestamp": unified_state.last_updated,
83 });
84
85 Ok(Json(summary))
86}
87
88pub async fn get_state(
92 State(state): State<XRayState>,
93 Query(params): Query<XRayQuery>,
94) -> Result<Json<Value>, StatusCode> {
95 let unified_state = state.engine.get_state(¶ms.workspace).await.ok_or_else(|| {
96 debug!("No state found for workspace: {}", params.workspace);
97 StatusCode::NOT_FOUND
98 })?;
99
100 Ok(Json(
101 serde_json::to_value(&unified_state).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
102 ))
103}
104
105pub async fn get_request_context(
112 State(state): State<XRayState>,
113 Path(request_id): Path<String>,
114 Query(params): Query<XRayQuery>,
115) -> Result<Json<Value>, StatusCode> {
116 let contexts = state.request_contexts.read().await;
118 if let Some(snapshot) = contexts.get(&request_id) {
119 if snapshot.workspace_id == params.workspace {
121 return Ok(Json(serde_json::json!({
122 "request_id": request_id,
123 "workspace": snapshot.workspace_id,
124 "state_snapshot": snapshot.state_snapshot,
125 "timestamp": snapshot.timestamp,
126 "cached": true,
127 })));
128 }
129 }
130 drop(contexts);
131
132 debug!(
134 "Request context not found for request_id: {}, returning current state",
135 request_id
136 );
137 let unified_state = state.engine.get_state(¶ms.workspace).await.ok_or_else(|| {
138 debug!("No state found for workspace: {}", params.workspace);
139 StatusCode::NOT_FOUND
140 })?;
141
142 Ok(Json(serde_json::json!({
143 "request_id": request_id,
144 "workspace": params.workspace,
145 "state_snapshot": serde_json::to_value(&unified_state).unwrap_or_default(),
146 "timestamp": unified_state.last_updated,
147 "cached": false,
148 "note": "Snapshot not found, returning current state",
149 })))
150}
151
152pub async fn store_request_context(
157 state: &XRayState,
158 request_id: String,
159 workspace_id: String,
160 unified_state: &mockforge_core::consistency::types::UnifiedState,
161) {
162 let state_snapshot = serde_json::to_value(unified_state).unwrap_or_default();
163 let snapshot = RequestContextSnapshot {
164 workspace_id: workspace_id.clone(),
165 state_snapshot,
166 timestamp: unified_state.last_updated.timestamp(),
167 };
168
169 let mut contexts = state.request_contexts.write().await;
170
171 let workspace_entries: Vec<_> = contexts
174 .iter()
175 .filter(|(_, s)| s.workspace_id == workspace_id)
176 .map(|(k, _)| k.clone())
177 .collect();
178
179 if workspace_entries.len() >= 1000 {
180 let mut timestamps: Vec<_> = workspace_entries
182 .iter()
183 .filter_map(|id| contexts.get(id).map(|s| (id.clone(), s.timestamp)))
184 .collect();
185 timestamps.sort_by_key(|(_, ts)| *ts);
186
187 for (id, _) in timestamps.iter().take(100) {
188 contexts.remove(id);
189 }
190 }
191
192 contexts.insert(request_id, snapshot);
193}
194
195pub async fn get_workspace_summary(
199 State(state): State<XRayState>,
200 Path(workspace_id): Path<String>,
201) -> Result<Json<Value>, StatusCode> {
202 let unified_state = state.engine.get_state(&workspace_id).await.ok_or_else(|| {
203 debug!("No state found for workspace: {}", workspace_id);
204 StatusCode::NOT_FOUND
205 })?;
206
207 let summary = serde_json::json!({
208 "workspace_id": unified_state.workspace_id,
209 "scenario": unified_state.active_scenario,
210 "persona_id": unified_state.active_persona.as_ref().map(|p| p.id.clone()),
211 "reality_level": unified_state.reality_level.value(),
212 "reality_ratio": unified_state.reality_continuum_ratio,
213 "active_chaos_rules_count": unified_state.active_chaos_rules.len(),
214 "entity_count": unified_state.entity_state.len(),
215 "protocol_count": unified_state.protocol_states.len(),
216 "last_updated": unified_state.last_updated,
217 });
218
219 Ok(Json(summary))
220}
221
222pub async fn list_entities(
226 State(state): State<XRayState>,
227 Query(params): Query<XRayQuery>,
228) -> Result<Json<Value>, StatusCode> {
229 let unified_state = state.engine.get_state(¶ms.workspace).await.ok_or_else(|| {
230 debug!("No state found for workspace: {}", params.workspace);
231 StatusCode::NOT_FOUND
232 })?;
233
234 let entities: Vec<&mockforge_core::consistency::EntityState> =
235 unified_state.entity_state.values().collect();
236
237 Ok(Json(serde_json::json!({
238 "workspace": params.workspace,
239 "entities": entities,
240 "count": entities.len(),
241 })))
242}
243
244pub async fn get_entity(
248 State(state): State<XRayState>,
249 Path((entity_type, entity_id)): Path<(String, String)>,
250 Query(params): Query<XRayQuery>,
251) -> Result<Json<Value>, StatusCode> {
252 let entity = state
253 .engine
254 .get_entity(¶ms.workspace, &entity_type, &entity_id)
255 .await
256 .ok_or_else(|| {
257 debug!(
258 "Entity not found: {}:{} in workspace: {}",
259 entity_type, entity_id, params.workspace
260 );
261 StatusCode::NOT_FOUND
262 })?;
263
264 Ok(Json(
265 serde_json::to_value(&entity).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
266 ))
267}
268
269pub fn xray_router(state: XRayState) -> axum::Router {
271 use axum::routing::get;
272
273 axum::Router::new()
274 .route("/api/v1/xray/state/summary", get(get_state_summary))
275 .route("/api/v1/xray/state", get(get_state))
276 .route("/api/v1/xray/request-context/{request_id}", get(get_request_context))
277 .route("/api/v1/xray/workspace/{workspace_id}/summary", get(get_workspace_summary))
278 .route("/api/v1/xray/entities", get(list_entities))
279 .route("/api/v1/xray/entities/{entity_type}/{entity_id}", get(get_entity))
280 .with_state(state)
281}