systemprompt-cli 0.2.1

mod ai_artifacts;
mod ai_display;
mod ai_mcp;
mod ai_trace_display;
mod display;
mod json;
mod list;
mod show;
mod summary;

pub use summary::{SummaryContext, print_summary};

use anyhow::Result;
use clap::Subcommand;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use systemprompt_identifiers::{AiRequestId, TraceId};
use systemprompt_runtime::DatabaseContext;

use super::types::{MessageRow, ToolCallRow};
use crate::CliConfig;
use crate::shared::render_result;

#[derive(Debug, Subcommand)]
pub enum TraceCommands {
    #[command(
        about = "List recent traces",
        after_help = "EXAMPLES:\n  systemprompt infra logs trace list\n  systemprompt infra logs \
                      trace list --limit 50 --since 1h\n  systemprompt infra logs trace list \
                      --agent researcher --status completed"
    )]
    List(list::ListArgs),

    #[command(
        about = "Show trace details",
        after_help = "EXAMPLES:\n  systemprompt infra logs trace show abc123\n  systemprompt \
                      infra logs trace show abc123 --verbose\n  systemprompt infra logs trace \
                      show abc123 --steps --ai --mcp"
    )]
    Show(show::ShowArgs),
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct TraceEventRow {
    pub timestamp: String,
    pub delta_ms: i64,
    pub event_type: String,
    pub details: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub latency_ms: Option<i64>,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, JsonSchema)]
pub struct AiSummaryRow {
    pub request_count: i64,
    pub total_tokens: i64,
    pub input_tokens: i64,
    pub output_tokens: i64,
    pub cost_dollars: f64,
    pub total_latency_ms: i64,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, JsonSchema)]
pub struct McpSummaryRow {
    pub execution_count: i64,
    pub total_execution_time_ms: i64,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, JsonSchema)]
pub struct StepSummaryRow {
    pub total: i64,
    pub completed: i64,
    pub failed: i64,
    pub pending: i64,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct TraceViewOutput {
    pub trace_id: TraceId,
    pub events: Vec<TraceEventRow>,
    pub ai_summary: AiSummaryRow,
    pub mcp_summary: McpSummaryRow,
    pub step_summary: StepSummaryRow,
    #[serde(skip_serializing_if = "Option::is_none", rename = "task_id")]
    pub task: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duration_ms: Option<i64>,
    pub status: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct TraceListRow {
    pub trace_id: TraceId,
    pub timestamp: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent: Option<String>,
    pub status: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duration_ms: Option<i64>,
    pub ai_requests: i64,
    pub mcp_calls: i64,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct TraceListOutput {
    pub traces: Vec<TraceListRow>,
    pub total: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct TaskInfoRow {
    #[serde(rename = "task_id")]
    pub task: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent_name: Option<String>,
    pub status: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub started_at: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duration_ms: Option<i64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct StepRow {
    pub step_number: i32,
    pub step_type: String,
    pub title: String,
    pub status: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duration_ms: Option<i64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct AiRequestRow {
    pub request_id: AiRequestId,
    pub model: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<i32>,
    pub tokens: String,
    pub cost: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub latency_ms: Option<i64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct ArtifactRow {
    #[serde(rename = "artifact_id")]
    pub artifact: String,
    pub artifact_type: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub source: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_name: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct AiTraceOutput {
    #[serde(rename = "task_id")]
    pub task: String,
    pub task_info: TaskInfoRow,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user_input: Option<String>,
    pub execution_steps: Vec<StepRow>,
    pub ai_requests: Vec<AiRequestRow>,
    pub mcp_executions: Vec<ToolCallRow>,
    pub artifacts: Vec<ArtifactRow>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent_response: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct AiLookupOutput {
    pub request_id: AiRequestId,
    pub provider: String,
    pub model: String,
    pub input_tokens: i32,
    pub output_tokens: i32,
    pub cost_dollars: f64,
    pub latency_ms: i64,
    pub messages: Vec<MessageRow>,
    pub linked_mcp_calls: Vec<ToolCallRow>,
}

pub async fn execute(command: TraceCommands, config: &CliConfig) -> Result<()> {
    match command {
        TraceCommands::List(args) => list::execute(args, config).await,
        TraceCommands::Show(args) => {
            let result = show::execute(args).await?;
            render_result(&result);
            Ok(())
        },
    }
}

pub async fn execute_with_pool(
    command: TraceCommands,
    db_ctx: &DatabaseContext,
    config: &CliConfig,
) -> Result<()> {
    match command {
        TraceCommands::List(args) => list::execute_with_pool(args, db_ctx, config).await,
        TraceCommands::Show(args) => {
            let result = show::execute_with_pool(args, db_ctx).await?;
            render_result(&result);
            Ok(())
        },
    }
}