klieo-ops-api 2.2.0

use axum::{
    extract::{Path, Query, State},
    Json,
};
use klieo_core::ids::RunId;
use klieo_runlog::types::RunStatus;
use serde::Deserialize;
use std::sync::Arc;

use crate::error::OpsError;
use crate::state::OpsState;
use crate::types::{AgentSummary, RunLogSummary};

const DEFAULT_LIMIT: usize = 50;
const MAX_LIMIT: usize = 200;
/// Production stores MUST handle this as a simple row limit, not an unbounded query.
const AGGREGATION_SCAN_LIMIT: usize = 10_000;

#[derive(Debug, Deserialize)]
pub(crate) struct ListRunsParams {
    pub(crate) agent: Option<String>,
    pub(crate) status: Option<String>,
    pub(crate) limit: Option<usize>,
}

pub(crate) async fn list_runs(
    State(state): State<Arc<OpsState>>,
    Query(params): Query<ListRunsParams>,
) -> Result<Json<Vec<RunLogSummary>>, OpsError> {
    let store = state
        .run_log_store
        .as_ref()
        .ok_or(OpsError::StoreNotConfigured)?;

    let limit = params.limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT);
    let mut logs = store
        .list(params.agent.as_deref(), limit)
        .await
        .map_err(|e| {
            tracing::error!(operation = "list_runs", detail = %e, "store error");
            OpsError::Internal(Box::new(e))
        })?;

    if let Some(status_str) = &params.status {
        let target = parse_run_status(status_str)?;
        logs.retain(|l| std::mem::discriminant(&l.status) == std::mem::discriminant(&target));
    }

    Ok(Json(logs.iter().map(RunLogSummary::from_run_log).collect()))
}

fn parse_run_status(s: &str) -> Result<RunStatus, OpsError> {
    match s {
        "running" => Ok(RunStatus::Running),
        "completed" => Ok(RunStatus::Completed),
        "failed" => Ok(RunStatus::Failed),
        "cancelled" => Ok(RunStatus::Cancelled),
        _ => Err(OpsError::InvalidFilter),
    }
}

pub(crate) async fn get_run(
    State(state): State<Arc<OpsState>>,
    Path(run_id_str): Path<String>,
) -> Result<Json<klieo_runlog::RunLog>, OpsError> {
    let store = state
        .run_log_store
        .as_ref()
        .ok_or(OpsError::StoreNotConfigured)?;

    let run_id = ulid::Ulid::from_string(&run_id_str)
        .map(RunId)
        .map_err(|_| OpsError::RunNotFound)?;

    let log = store
        .get(run_id)
        .await
        .map_err(|e| {
            tracing::error!(operation = "get_run", run_id = %run_id_str, detail = %e, "store error");
            OpsError::Internal(Box::new(e))
        })?
        .ok_or(OpsError::RunNotFound)?;

    Ok(Json(log))
}

pub(crate) async fn list_agents(
    State(state): State<Arc<OpsState>>,
) -> Result<Json<Vec<AgentSummary>>, OpsError> {
    let store = state
        .run_log_store
        .as_ref()
        .ok_or(OpsError::StoreNotConfigured)?;

    let logs = store
        .list(None, AGGREGATION_SCAN_LIMIT)
        .await
        .map_err(|e| {
            tracing::error!(operation = "list_agents", detail = %e, "store error");
            OpsError::Internal(Box::new(e))
        })?;

    let mut map: std::collections::HashMap<String, AgentSummary> = std::collections::HashMap::new();

    for log in &logs {
        let entry = map
            .entry(log.agent.clone())
            .or_insert_with(|| AgentSummary {
                agent: log.agent.clone(),
                last_seen: log.started_at,
                runs_running: 0,
                runs_completed: 0,
                runs_failed: 0,
                runs_cancelled: 0,
            });
        if log.started_at > entry.last_seen {
            entry.last_seen = log.started_at;
        }
        match log.status {
            RunStatus::Running => entry.runs_running += 1,
            RunStatus::Completed => entry.runs_completed += 1,
            RunStatus::Failed => entry.runs_failed += 1,
            RunStatus::Cancelled => entry.runs_cancelled += 1,
            _ => {}
        }
    }

    let mut summaries: Vec<AgentSummary> = map.into_values().collect();
    summaries.sort_by(|a, b| b.last_seen.cmp(&a.last_seen));
    Ok(Json(summaries))
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{
        body::Body,
        http::{Request, StatusCode},
        routing::get,
        Router,
    };
    use chrono::Utc;
    use klieo_auth_common::AllowAnonymous;
    use klieo_runlog::{
        types::{RunStatus, Usage},
        InMemoryRunLogStore, RunLog, RunLogStore,
    };
    use std::sync::Arc;
    use tower::ServiceExt;

    fn make_state_with_store(store: Arc<dyn klieo_runlog::RunLogStore>) -> Arc<OpsState> {
        Arc::new(OpsState {
            authenticator: Arc::new(AllowAnonymous),
            run_log_store: Some(store),
            task_store: None,
            resume_buffer: None,
        })
    }

    fn make_run(agent: &str, status: RunStatus) -> RunLog {
        RunLog {
            run_id: RunId::new(),
            agent: agent.to_string(),
            started_at: Utc::now(),
            finished_at: None,
            status,
            steps: vec![],
            tokens: Usage::default(),
            cost_estimate: None,
        }
    }

    fn ops_router(state: Arc<OpsState>) -> Router {
        Router::new()
            .route("/runs", get(list_runs))
            .route("/runs/:run_id", get(get_run))
            .route("/agents", get(list_agents))
            .with_state(state)
    }

    #[tokio::test]
    async fn list_runs_empty_store_returns_empty_array() {
        let store = Arc::new(InMemoryRunLogStore::new());
        let state = make_state_with_store(store);
        let app = ops_router(state);

        let req = Request::builder().uri("/runs").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json, serde_json::json!([]));
    }

    #[tokio::test]
    async fn list_runs_no_store_returns_501() {
        let state = Arc::new(OpsState {
            authenticator: Arc::new(AllowAnonymous),
            run_log_store: None,
            task_store: None,
            resume_buffer: None,
        });
        let app = ops_router(state);
        let req = Request::builder().uri("/runs").body(Body::empty()).unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_IMPLEMENTED);
    }

    #[tokio::test]
    async fn list_runs_filters_by_status() {
        let store = Arc::new(InMemoryRunLogStore::new());
        let run_a = make_run("agent-a", RunStatus::Running);
        let run_b = make_run("agent-b", RunStatus::Completed);
        store.put(&run_a).await.unwrap();
        store.put(&run_b).await.unwrap();

        let state = make_state_with_store(Arc::clone(&store) as Arc<dyn klieo_runlog::RunLogStore>);
        let app = ops_router(state);
        let req = Request::builder()
            .uri("/runs?status=running")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: Vec<serde_json::Value> = serde_json::from_slice(&body).unwrap();
        assert_eq!(json.len(), 1);
        assert_eq!(json[0]["agent"], "agent-a");
    }

    #[tokio::test]
    async fn get_run_returns_501_when_no_store() {
        let state = Arc::new(OpsState {
            authenticator: Arc::new(AllowAnonymous),
            run_log_store: None,
            task_store: None,
            resume_buffer: None,
        });
        let app = ops_router(state);
        let req = Request::builder()
            .uri(format!("/runs/{}", RunId::new()))
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_IMPLEMENTED);
    }

    #[tokio::test]
    async fn get_run_returns_404_for_unknown_id() {
        let store = Arc::new(InMemoryRunLogStore::new());
        let state = make_state_with_store(store);
        let app = ops_router(state);
        let req = Request::builder()
            .uri(format!("/runs/{}", RunId::new()))
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test]
    async fn get_run_returns_full_log() {
        let store = Arc::new(InMemoryRunLogStore::new());
        let run = make_run("writer", RunStatus::Completed);
        let run_id_str = run.run_id.to_string();
        store.put(&run).await.unwrap();

        let state = make_state_with_store(Arc::clone(&store) as Arc<dyn klieo_runlog::RunLogStore>);
        let app = ops_router(state);
        let req = Request::builder()
            .uri(format!("/runs/{run_id_str}"))
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["agent"], "writer");
    }

    #[tokio::test]
    async fn list_runs_unknown_status_returns_400() {
        let store = Arc::new(InMemoryRunLogStore::new());
        let state = make_state_with_store(store);
        let app = ops_router(state);
        let req = Request::builder()
            .uri("/runs?status=garbage")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }

    #[tokio::test]
    async fn list_agents_no_store_returns_501() {
        let state = Arc::new(OpsState {
            authenticator: Arc::new(AllowAnonymous),
            run_log_store: None,
            task_store: None,
            resume_buffer: None,
        });
        let app = ops_router(state);
        let req = Request::builder()
            .uri("/agents")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_IMPLEMENTED);
    }

    #[tokio::test]
    async fn get_run_returns_404_for_malformed_id() {
        let store = Arc::new(InMemoryRunLogStore::new());
        let state = make_state_with_store(store);
        let app = ops_router(state);
        let req = Request::builder()
            .uri("/runs/not-a-valid-ulid")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test]
    async fn list_agents_aggregates_counts_by_status() {
        let store = Arc::new(InMemoryRunLogStore::new());
        store
            .put(&make_run("alpha", RunStatus::Running))
            .await
            .unwrap();
        store
            .put(&make_run("alpha", RunStatus::Completed))
            .await
            .unwrap();
        store
            .put(&make_run("beta", RunStatus::Failed))
            .await
            .unwrap();

        let state = make_state_with_store(Arc::clone(&store) as Arc<dyn klieo_runlog::RunLogStore>);
        let app = ops_router(state);
        let req = Request::builder()
            .uri("/agents")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let agents: Vec<serde_json::Value> = serde_json::from_slice(&body).unwrap();
        let alpha = agents.iter().find(|a| a["agent"] == "alpha").unwrap();
        assert_eq!(alpha["runs_running"], 1);
        assert_eq!(alpha["runs_completed"], 1);
        let beta = agents.iter().find(|a| a["agent"] == "beta").unwrap();
        assert_eq!(beta["runs_failed"], 1);
    }
}