agent-infra-cli 0.4.0

use agent_infra_sdk::{
    InfraClient,
    trace_contract::{
        CreateLangfuseProjectRequest, InitializeRuntimeBindingRequest, ListLangfuseProjectsRequest,
        MAX_METRICS_LOOKBACK_DAYS, MetricsDataSource, MetricsQuery,
    },
};
use anyhow::Result;
use clap::{Args, Subcommand, ValueEnum};

use crate::{CommandOutput, JsonInput, parse_agent_id};

#[derive(Debug, Subcommand)]
pub enum TraceCommand {
    /// Query Langfuse metrics through Trace Infra.
    Metrics {
        #[command(subcommand)]
        command: TraceMetricsCommand,
    },
    /// Allocate and query Trace-owned Langfuse projects.
    LangfuseProject {
        #[command(subcommand)]
        command: LangfuseProjectCommand,
    },
    /// Inspect or initialize a runtime's Langfuse binding.
    Binding {
        #[command(subcommand)]
        command: TraceBindingCommand,
    },
}

#[derive(Debug, Subcommand)]
pub enum LangfuseProjectCommand {
    /// Create a project from a CreateLangfuseProjectRequest JSON document.
    Create {
        #[command(flatten)]
        input: JsonInput,
        #[arg(long)]
        idempotency_key: String,
    },
    Get {
        project_id: String,
    },
    List(LangfuseProjectListArgs),
}

#[derive(Debug, Args)]
pub struct LangfuseProjectListArgs {
    #[arg(long)]
    pub(crate) related_agent_id: Option<String>,
    #[arg(long)]
    pub(crate) related_runtime_id: Option<String>,
    #[arg(long)]
    pub(crate) related_user_id: Option<String>,
    #[arg(long)]
    pub(crate) cursor: Option<String>,
    #[arg(long, value_parser = clap::value_parser!(u32).range(1..=200))]
    pub(crate) limit: Option<u32>,
}

#[derive(Debug, Subcommand)]
pub enum TraceBindingCommand {
    /// Get the Langfuse project bound to a runtime in the current tenant/project.
    Get { runtime_id: String },
    /// Allocate or replay the Langfuse binding for a runtime.
    Initialize {
        runtime_id: String,
        #[arg(long)]
        project_name: Option<String>,
        #[arg(long)]
        agent_id: Option<String>,
        #[arg(long)]
        owner_email: Option<String>,
        #[arg(long)]
        idempotency_key: Option<String>,
    },
}

#[derive(Debug, Subcommand)]
pub enum TraceMetricsCommand {
    TokenUsage(TraceMetricsArgs),
    ModelBreakdown(TraceMetricsArgs),
    ProjectSummary(TraceMetricsArgs),
    Latency(TraceMetricsArgs),
    ErrorRate(TraceMetricsArgs),
    CacheHitRate(TraceMetricsArgs),
    CostBreakdown(TraceMetricsArgs),
    ModelTrend(TraceMetricsArgs),
    TokenDistribution(TraceMetricsArgs),
    Activity(TraceMetricsArgs),
    PricingTier(TraceMetricsArgs),
}

#[derive(Debug, Args)]
pub struct TraceMetricsArgs {
    /// Filter by one or more Langfuse project IDs.
    #[arg(long = "project-id")]
    project_ids: Vec<String>,
    /// Number of days to include in the query window.
    #[arg(long, default_value_t = 30, value_parser = clap::value_parser!(u32).range(1..=3650))]
    days: u32,
    /// Select the Langfuse v3/v4 ClickHouse data path.
    #[arg(long, value_enum, default_value_t = CliMetricsDataSource::Auto)]
    data_source: CliMetricsDataSource,
    /// v3/v4 cutover boundary used when --data-source is auto.
    #[arg(long, default_value = "2026-09-01 00:00:00")]
    v4_cutover_date: String,
}

impl TraceMetricsArgs {
    pub(crate) fn into_query(self) -> MetricsQuery {
        MetricsQuery {
            project_ids: (!self.project_ids.is_empty()).then_some(self.project_ids),
            days: self.days.min(MAX_METRICS_LOOKBACK_DAYS),
            data_source: self.data_source.into(),
            v4_cutover_date: self.v4_cutover_date,
        }
    }
}

#[derive(Clone, Copy, Debug, Default, ValueEnum)]
pub enum CliMetricsDataSource {
    V3,
    V4,
    #[default]
    Auto,
}

impl From<CliMetricsDataSource> for MetricsDataSource {
    fn from(value: CliMetricsDataSource) -> Self {
        match value {
            CliMetricsDataSource::V3 => Self::V3,
            CliMetricsDataSource::V4 => Self::V4,
            CliMetricsDataSource::Auto => Self::Auto,
        }
    }
}

pub(crate) async fn execute(client: &InfraClient, command: TraceCommand) -> Result<CommandOutput> {
    match command {
        TraceCommand::Metrics { command } => execute_metrics(client, command).await,
        TraceCommand::LangfuseProject { command } => match command {
            LangfuseProjectCommand::Create {
                input,
                idempotency_key,
            } => {
                let request: CreateLangfuseProjectRequest = input.read()?;
                CommandOutput::json(
                    client
                        .admin()
                        .create_langfuse_project(&request, &idempotency_key)
                        .await?,
                )
            }
            LangfuseProjectCommand::Get { project_id } => {
                CommandOutput::json(client.admin().langfuse_project(&project_id).await?)
            }
            LangfuseProjectCommand::List(args) => CommandOutput::json(
                client
                    .admin()
                    .list_langfuse_projects(&ListLangfuseProjectsRequest {
                        related_agent_id: args
                            .related_agent_id
                            .as_deref()
                            .map(parse_agent_id)
                            .transpose()?,
                        related_runtime_id: args.related_runtime_id,
                        related_user_id: args.related_user_id,
                        cursor: args.cursor,
                        limit: args.limit,
                    })
                    .await?,
            ),
        },
        TraceCommand::Binding { command } => match command {
            TraceBindingCommand::Get { runtime_id } => {
                CommandOutput::json(client.trace().get_runtime_binding(&runtime_id).await?)
            }
            TraceBindingCommand::Initialize {
                runtime_id,
                project_name,
                agent_id,
                owner_email,
                idempotency_key,
            } => {
                let idempotency_key =
                    idempotency_key.unwrap_or_else(|| format!("init-trace-{runtime_id}"));
                CommandOutput::json(
                    client
                        .trace()
                        .initialize_runtime_binding(
                            &runtime_id,
                            &InitializeRuntimeBindingRequest {
                                project_name,
                                agent_id: agent_id.as_deref().map(parse_agent_id).transpose()?,
                                owner_email,
                            },
                            &idempotency_key,
                        )
                        .await?,
                )
            }
        },
    }
}

async fn execute_metrics(
    client: &InfraClient,
    command: TraceMetricsCommand,
) -> Result<CommandOutput> {
    match command {
        TraceMetricsCommand::TokenUsage(args) => {
            CommandOutput::json(client.trace().token_usage(&args.into_query()).await?)
        }
        TraceMetricsCommand::ModelBreakdown(args) => {
            CommandOutput::json(client.trace().model_breakdown(&args.into_query()).await?)
        }
        TraceMetricsCommand::ProjectSummary(args) => {
            CommandOutput::json(client.trace().project_summary(&args.into_query()).await?)
        }
        TraceMetricsCommand::Latency(args) => {
            CommandOutput::json(client.trace().latency(&args.into_query()).await?)
        }
        TraceMetricsCommand::ErrorRate(args) => {
            CommandOutput::json(client.trace().error_rate(&args.into_query()).await?)
        }
        TraceMetricsCommand::CacheHitRate(args) => {
            CommandOutput::json(client.trace().cache_hit_rate(&args.into_query()).await?)
        }
        TraceMetricsCommand::CostBreakdown(args) => {
            CommandOutput::json(client.trace().cost_breakdown(&args.into_query()).await?)
        }
        TraceMetricsCommand::ModelTrend(args) => {
            CommandOutput::json(client.trace().model_trend(&args.into_query()).await?)
        }
        TraceMetricsCommand::TokenDistribution(args) => CommandOutput::json(
            client
                .trace()
                .token_distribution(&args.into_query())
                .await?,
        ),
        TraceMetricsCommand::Activity(args) => {
            CommandOutput::json(client.trace().activity(&args.into_query()).await?)
        }
        TraceMetricsCommand::PricingTier(args) => {
            CommandOutput::json(client.trace().pricing_tier(&args.into_query()).await?)
        }
    }
}