agent-infra-cli 0.6.0

//! `infra execution` — durable run and event operations.

use agent_infra_sdk::InfraClient;
use agent_infra_sdk::execution_contract::{
    AppendExecutionEventRequest, CreateExecutionRunRequest, GetExecutionEventsRequest,
    GetExecutionRunRequest, ListExecutionRunsRequest, UpdateExecutionRunRequest,
};
use anyhow::Result;
use clap::{Args, Subcommand, ValueEnum};

use crate::{CommandOutput, JsonInput};

#[derive(Debug, Subcommand)]
pub enum ExecutionCommand {
    /// Create a run from a CreateExecutionRunRequest JSON document.
    Create {
        #[command(flatten)]
        input: JsonInput,
    },
    /// Update a run from an UpdateExecutionRunRequest JSON document.
    Update {
        #[command(flatten)]
        input: JsonInput,
    },
    /// Get one run by ID or idempotency key.
    Get {
        #[arg(long, conflicts_with = "idempotency_key")]
        run_id: Option<String>,
        #[arg(long, conflicts_with = "run_id")]
        idempotency_key: Option<String>,
    },
    /// List runs by status, session, or limit.
    List(ExecutionListArgs),
    /// Cancel a queued or running run.
    Cancel { run_id: String },
    /// Append an execution event from an AppendExecutionEventRequest JSON document.
    Event {
        run_id: String,
        #[command(flatten)]
        input: JsonInput,
    },
    /// List execution events for a session.
    Events {
        session_id: String,
        #[arg(long, value_parser = clap::value_parser!(u64).range(1..=1000))]
        limit: Option<u64>,
    },
}

#[derive(Debug, Args)]
pub struct ExecutionListArgs {
    #[arg(long, value_enum)]
    status: Option<CliExecutionStatus>,
    #[arg(long)]
    session_id: Option<String>,
    #[arg(long, value_parser = clap::value_parser!(u64).range(1..=1000))]
    limit: Option<u64>,
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum CliExecutionStatus {
    Queued,
    Running,
    Completed,
    Failed,
    Canceled,
}

impl From<CliExecutionStatus> for agent_infra_sdk::execution_contract::ExecutionRunStatus {
    fn from(value: CliExecutionStatus) -> Self {
        match value {
            CliExecutionStatus::Queued => Self::Queued,
            CliExecutionStatus::Running => Self::Running,
            CliExecutionStatus::Completed => Self::Completed,
            CliExecutionStatus::Failed => Self::Failed,
            CliExecutionStatus::Canceled => Self::Canceled,
        }
    }
}

pub(crate) async fn execute(
    client: &InfraClient,
    command: ExecutionCommand,
) -> Result<CommandOutput> {
    let execution = client.execution();
    match command {
        ExecutionCommand::Create { input } => {
            let request: CreateExecutionRunRequest = input.read()?;
            CommandOutput::json(execution.create_run(&request).await?)
        }
        ExecutionCommand::Update { input } => {
            let request: UpdateExecutionRunRequest = input.read()?;
            CommandOutput::json(execution.update_run(&request).await?)
        }
        ExecutionCommand::Get {
            run_id,
            idempotency_key,
        } => {
            let request = GetExecutionRunRequest {
                run_id,
                idempotency_key,
            };
            CommandOutput::json(execution.get_run(&request).await?)
        }
        ExecutionCommand::List(args) => {
            let request = ListExecutionRunsRequest {
                status: args.status.map(Into::into),
                session_id: args.session_id,
                limit: args.limit.map(|limit| limit as usize),
            };
            CommandOutput::json(execution.runs(&request).await?)
        }
        ExecutionCommand::Cancel { run_id } => {
            CommandOutput::json(execution.cancel_run(&run_id).await?)
        }
        ExecutionCommand::Event { run_id, input } => {
            let request: AppendExecutionEventRequest = input.read()?;
            CommandOutput::json(execution.append_event(&run_id, &request).await?)
        }
        ExecutionCommand::Events { session_id, limit } => {
            let request = GetExecutionEventsRequest {
                session_id,
                limit: limit.map(|limit| limit as usize),
            };
            CommandOutput::json(execution.events(&request).await?)
        }
    }
}