use axum::{
extract::{Path, Query, State},
http::StatusCode,
response::Json,
};
use mockforge_core::consistency::ConsistencyEngine;
use serde::Deserialize;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::debug;
#[derive(Debug, Clone)]
pub(crate) struct RequestContextSnapshot {
workspace_id: String,
state_snapshot: Value,
timestamp: i64,
}
#[derive(Clone)]
pub struct XRayState {
pub engine: Arc<ConsistencyEngine>,
pub(crate) request_contexts: Arc<RwLock<HashMap<String, RequestContextSnapshot>>>,
}
#[derive(Debug, Deserialize)]
pub struct XRayQuery {
#[serde(default = "default_workspace")]
pub workspace: String,
}
fn default_workspace() -> String {
"default".to_string()
}
pub async fn get_state_summary(
State(state): State<XRayState>,
Query(params): Query<XRayQuery>,
) -> Result<Json<Value>, StatusCode> {
let unified_state = state.engine.get_state(¶ms.workspace).await.ok_or_else(|| {
debug!("No state found for workspace: {}", params.workspace);
StatusCode::NOT_FOUND
})?;
let summary = serde_json::json!({
"workspace_id": unified_state.workspace_id,
"scenario": unified_state.active_scenario,
"persona": unified_state.active_persona.as_ref().map(|p| serde_json::json!({
"id": p.id,
"traits": p.traits,
})),
"reality_level": unified_state.reality_level.value(),
"reality_level_name": unified_state.reality_level.name(),
"reality_ratio": unified_state.reality_continuum_ratio,
"chaos_rules": unified_state
.active_chaos_rules
.iter()
.filter_map(|r| r.get("name").and_then(|v| v.as_str()).map(|s| s.to_string()))
.collect::<Vec<_>>(),
"timestamp": unified_state.last_updated,
});
Ok(Json(summary))
}
pub async fn get_state(
State(state): State<XRayState>,
Query(params): Query<XRayQuery>,
) -> Result<Json<Value>, StatusCode> {
let unified_state = state.engine.get_state(¶ms.workspace).await.ok_or_else(|| {
debug!("No state found for workspace: {}", params.workspace);
StatusCode::NOT_FOUND
})?;
Ok(Json(
serde_json::to_value(&unified_state).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
))
}
pub async fn get_request_context(
State(state): State<XRayState>,
Path(request_id): Path<String>,
Query(params): Query<XRayQuery>,
) -> Result<Json<Value>, StatusCode> {
let contexts = state.request_contexts.read().await;
if let Some(snapshot) = contexts.get(&request_id) {
if snapshot.workspace_id == params.workspace {
return Ok(Json(serde_json::json!({
"request_id": request_id,
"workspace": snapshot.workspace_id,
"state_snapshot": snapshot.state_snapshot,
"timestamp": snapshot.timestamp,
"cached": true,
})));
}
}
drop(contexts);
debug!(
"Request context not found for request_id: {}, returning current state",
request_id
);
let unified_state = state.engine.get_state(¶ms.workspace).await.ok_or_else(|| {
debug!("No state found for workspace: {}", params.workspace);
StatusCode::NOT_FOUND
})?;
Ok(Json(serde_json::json!({
"request_id": request_id,
"workspace": params.workspace,
"state_snapshot": serde_json::to_value(&unified_state).unwrap_or_default(),
"timestamp": unified_state.last_updated,
"cached": false,
"note": "Snapshot not found, returning current state",
})))
}
pub async fn store_request_context(
state: &XRayState,
request_id: String,
workspace_id: String,
unified_state: &mockforge_core::consistency::types::UnifiedState,
) {
let state_snapshot = serde_json::to_value(unified_state).unwrap_or_default();
let snapshot = RequestContextSnapshot {
workspace_id: workspace_id.clone(),
state_snapshot,
timestamp: unified_state.last_updated.timestamp(),
};
let mut contexts = state.request_contexts.write().await;
let workspace_entries: Vec<_> = contexts
.iter()
.filter(|(_, s)| s.workspace_id == workspace_id)
.map(|(k, _)| k.clone())
.collect();
if workspace_entries.len() >= 1000 {
let mut timestamps: Vec<_> = workspace_entries
.iter()
.filter_map(|id| contexts.get(id).map(|s| (id.clone(), s.timestamp)))
.collect();
timestamps.sort_by_key(|(_, ts)| *ts);
for (id, _) in timestamps.iter().take(100) {
contexts.remove(id);
}
}
contexts.insert(request_id, snapshot);
}
pub async fn get_workspace_summary(
State(state): State<XRayState>,
Path(workspace_id): Path<String>,
) -> Result<Json<Value>, StatusCode> {
let unified_state = state.engine.get_state(&workspace_id).await.ok_or_else(|| {
debug!("No state found for workspace: {}", workspace_id);
StatusCode::NOT_FOUND
})?;
let summary = serde_json::json!({
"workspace_id": unified_state.workspace_id,
"scenario": unified_state.active_scenario,
"persona_id": unified_state.active_persona.as_ref().map(|p| p.id.clone()),
"reality_level": unified_state.reality_level.value(),
"reality_ratio": unified_state.reality_continuum_ratio,
"active_chaos_rules_count": unified_state.active_chaos_rules.len(),
"entity_count": unified_state.entity_state.len(),
"protocol_count": unified_state.protocol_states.len(),
"last_updated": unified_state.last_updated,
});
Ok(Json(summary))
}
pub async fn list_entities(
State(state): State<XRayState>,
Query(params): Query<XRayQuery>,
) -> Result<Json<Value>, StatusCode> {
let unified_state = state.engine.get_state(¶ms.workspace).await.ok_or_else(|| {
debug!("No state found for workspace: {}", params.workspace);
StatusCode::NOT_FOUND
})?;
let entities: Vec<&mockforge_core::consistency::EntityState> =
unified_state.entity_state.values().collect();
Ok(Json(serde_json::json!({
"workspace": params.workspace,
"entities": entities,
"count": entities.len(),
})))
}
pub async fn get_entity(
State(state): State<XRayState>,
Path((entity_type, entity_id)): Path<(String, String)>,
Query(params): Query<XRayQuery>,
) -> Result<Json<Value>, StatusCode> {
let entity = state
.engine
.get_entity(¶ms.workspace, &entity_type, &entity_id)
.await
.ok_or_else(|| {
debug!(
"Entity not found: {}:{} in workspace: {}",
entity_type, entity_id, params.workspace
);
StatusCode::NOT_FOUND
})?;
Ok(Json(
serde_json::to_value(&entity).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
))
}
pub fn xray_router(state: XRayState) -> axum::Router {
use axum::routing::get;
axum::Router::new()
.route("/api/v1/xray/state/summary", get(get_state_summary))
.route("/api/v1/xray/state", get(get_state))
.route("/api/v1/xray/request-context/{request_id}", get(get_request_context))
.route("/api/v1/xray/workspace/{workspace_id}/summary", get(get_workspace_summary))
.route("/api/v1/xray/entities", get(list_entities))
.route("/api/v1/xray/entities/{entity_type}/{entity_id}", get(get_entity))
.with_state(state)
}