agent_infra_cli/
lib.rs

1use std::{
2    fs,
3    io::{self, Read},
4    path::{Path, PathBuf},
5    sync::Arc,
6    time::Duration,
7};
8
9use agent_infra_sdk::{
10    CallOptions, InfraClient, StaticCredentials,
11    agent_contract::{
12        AgentId, AgentStatus, CreateAgentRequest, ListAgentsRequest, UpdateAgentRequest,
13    },
14    context_contract::ConversationRecord,
15    runtime_identity_contract::{
16        CreateRuntimeIdentityRequest, IssueRuntimeCredentialRequest, ReplacePermissionsRequest,
17        RuntimeIdentityId, RuntimeIdentityKind, TransitionIdentityRequest,
18        UpdateRuntimeIdentityRequest,
19    },
20    workspace_contract::{
21        CreateComputerRequest, CreateSpaceRequest, ListEnvironmentResourcesQuery, OwnerKind,
22    },
23};
24use anyhow::{Context, Result, anyhow, bail};
25use clap::{Args, Parser, Subcommand, ValueEnum};
26use serde::{Serialize, de::DeserializeOwned};
27use serde_json::Value;
28
29mod config;
30
31use config::{ConfigStore, resolve_credential};
32
33const DEFAULT_GATEWAY: &str = "http://127.0.0.1:5200";
34
35#[derive(Debug, Parser)]
36#[command(
37    name = "infra",
38    version,
39    about = "Operate Agent Infra through the typed Infra SDK",
40    arg_required_else_help = true
41)]
42pub struct Cli {
43    /// Agent Infra Gateway base URL.
44    #[arg(long, global = true)]
45    pub gateway: Option<String>,
46
47    /// Gateway bearer token. Prefer SKY_GATEWAY_TOKEN over command-line use.
48    #[arg(long, global = true, hide_env_values = true)]
49    pub token: Option<String>,
50
51    /// Override the local Infra CLI config file.
52    #[arg(long, global = true, env = "AGENT_INFRA_CONFIG", value_name = "PATH")]
53    pub config: Option<PathBuf>,
54
55    /// Override the Sky CLI home used for local credential discovery.
56    #[arg(
57        long,
58        global = true,
59        env = "AGENT_INFRA_SKY_HOME",
60        value_name = "PATH",
61        hide = true
62    )]
63    pub sky_home: Option<PathBuf>,
64
65    /// Request timeout in seconds.
66    #[arg(long, global = true, default_value_t = 30, value_parser = clap::value_parser!(u64).range(1..=600))]
67    pub timeout: u64,
68
69    /// Permit plaintext HTTP to non-loopback hosts protected by a trusted service mesh.
70    #[arg(long, global = true)]
71    pub trusted_mesh_http: bool,
72
73    /// JSON rendering mode.
74    #[arg(long, global = true, value_enum, default_value_t = OutputFormat::Pretty)]
75    pub output: OutputFormat,
76
77    #[command(subcommand)]
78    pub command: Command,
79}
80
81#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, ValueEnum)]
82pub enum OutputFormat {
83    /// Indented JSON.
84    #[default]
85    Pretty,
86    /// Compact JSON.
87    Json,
88}
89
90#[derive(Debug, Subcommand)]
91pub enum Command {
92    /// Persist and inspect local CLI settings.
93    Config {
94        #[command(subcommand)]
95        command: ConfigCommand,
96    },
97    /// Inspect local Gateway credential discovery.
98    Auth {
99        #[command(subcommand)]
100        command: AuthCommand,
101    },
102    /// Inspect Gateway health, discovery, routing, and policy.
103    Gateway {
104        #[command(subcommand)]
105        command: GatewayCommand,
106    },
107    /// Create and inspect Agent Registry resources.
108    Agent {
109        #[command(subcommand)]
110        command: AgentCommand,
111    },
112    /// Create and manage Runtime Identity resources.
113    Runtime {
114        #[command(subcommand)]
115        command: RuntimeCommand,
116    },
117    /// Read and append durable conversation context.
118    Context {
119        #[command(subcommand)]
120        command: ContextCommand,
121    },
122    /// Manage Environment computers.
123    Computer {
124        #[command(subcommand)]
125        command: ComputerCommand,
126    },
127    /// Manage Environment spaces.
128    Space {
129        #[command(subcommand)]
130        command: SpaceCommand,
131    },
132    /// Invoke the OpenAI-compatible Model Infra API.
133    Model {
134        #[command(subcommand)]
135        command: ModelCommand,
136    },
137}
138
139#[derive(Debug, Subcommand)]
140pub enum ConfigCommand {
141    /// Show the config file and effective Gateway selection.
142    Show,
143    /// Manage the persisted Gateway address.
144    Gateway {
145        #[command(subcommand)]
146        command: ConfigGatewayCommand,
147    },
148}
149
150#[derive(Debug, Subcommand)]
151pub enum ConfigGatewayCommand {
152    /// Persist a Gateway URL locally.
153    Set { url: String },
154    /// Show the persisted and effective Gateway URL.
155    Get,
156    /// Remove the persisted Gateway URL.
157    Unset,
158}
159
160#[derive(Debug, Subcommand)]
161pub enum AuthCommand {
162    /// Probe explicit, environment, and local Sky credential sources.
163    Status,
164}
165
166#[derive(Debug, Subcommand)]
167pub enum GatewayCommand {
168    Discovery,
169    Live,
170    Ready,
171    /// Poll until the Gateway and its dependencies are ready.
172    WaitReady {
173        #[arg(long, default_value_t = 30, value_parser = clap::value_parser!(u64).range(1..=3600))]
174        deadline: u64,
175    },
176    Dependencies,
177    Routes,
178    Metrics,
179    Policy,
180}
181
182#[derive(Debug, Subcommand)]
183pub enum AgentCommand {
184    /// Create an agent from a CreateAgentRequest JSON document.
185    Create(JsonInput),
186    Get {
187        agent_id: String,
188    },
189    List(AgentListArgs),
190    /// Update an agent from an UpdateAgentRequest JSON document.
191    Update {
192        agent_id: String,
193        #[command(flatten)]
194        input: JsonInput,
195    },
196}
197
198#[derive(Debug, Subcommand)]
199pub enum RuntimeCommand {
200    /// Create a named runtime identity.
201    Create(RuntimeCreateArgs),
202    /// Get a runtime identity by ID.
203    Get { runtime_id: String },
204    /// Change a runtime identity name without changing its ID.
205    Rename {
206        runtime_id: String,
207        name: String,
208        #[arg(long)]
209        expected_revision: u64,
210        #[arg(long)]
211        idempotency_key: Option<String>,
212    },
213    /// Activate a pending runtime identity using a TransitionIdentityRequest JSON document.
214    Activate {
215        runtime_id: String,
216        #[command(flatten)]
217        input: JsonInput,
218    },
219    /// Suspend an active runtime identity using a TransitionIdentityRequest JSON document.
220    Suspend {
221        runtime_id: String,
222        #[command(flatten)]
223        input: JsonInput,
224    },
225    /// Resume a suspended runtime identity using a TransitionIdentityRequest JSON document.
226    Resume {
227        runtime_id: String,
228        #[command(flatten)]
229        input: JsonInput,
230    },
231    /// Revoke a runtime identity using a TransitionIdentityRequest JSON document.
232    Revoke {
233        runtime_id: String,
234        #[command(flatten)]
235        input: JsonInput,
236    },
237    /// Replace the complete permission grant set from a JSON document.
238    Permissions {
239        runtime_id: String,
240        #[command(flatten)]
241        input: JsonInput,
242    },
243    /// Exchange a one-time bootstrap proof for the first workload credential.
244    BootstrapCredential {
245        #[command(flatten)]
246        input: JsonInput,
247    },
248}
249
250#[derive(Debug, Args)]
251pub struct RuntimeCreateArgs {
252    /// Mutable display name. Required unless --file supplies the full request.
253    #[arg(long, required_unless_present = "file")]
254    name: Option<String>,
255    /// Runtime workload kind. Defaults to code for inline creation.
256    #[arg(long, value_enum)]
257    kind: Option<CliRuntimeKind>,
258    /// Override the deterministic key derived from --name.
259    #[arg(long)]
260    idempotency_key: Option<String>,
261    /// Read a complete CreateRuntimeIdentityRequest JSON document.
262    #[arg(
263        long,
264        short = 'f',
265        value_name = "PATH",
266        conflicts_with_all = ["name", "kind", "idempotency_key"]
267    )]
268    file: Option<PathBuf>,
269}
270
271#[derive(Clone, Copy, Debug, ValueEnum)]
272enum CliRuntimeKind {
273    Code,
274    Evaluate,
275    Tool,
276    Custom,
277}
278
279impl From<CliRuntimeKind> for RuntimeIdentityKind {
280    fn from(value: CliRuntimeKind) -> Self {
281        match value {
282            CliRuntimeKind::Code => Self::Code,
283            CliRuntimeKind::Evaluate => Self::Evaluate,
284            CliRuntimeKind::Tool => Self::Tool,
285            CliRuntimeKind::Custom => Self::Custom,
286        }
287    }
288}
289
290#[derive(Debug, Args)]
291pub struct AgentListArgs {
292    #[arg(long, value_enum)]
293    status: Option<CliAgentStatus>,
294    #[arg(long)]
295    related_agent_id: Option<String>,
296    #[arg(long)]
297    related_runtime_id: Option<String>,
298    #[arg(long)]
299    related_user_id: Option<String>,
300    #[arg(long)]
301    cursor: Option<String>,
302    #[arg(long, value_parser = clap::value_parser!(u32).range(1..=200))]
303    limit: Option<u32>,
304}
305
306#[derive(Clone, Copy, Debug, ValueEnum)]
307enum CliAgentStatus {
308    Active,
309    Archived,
310}
311
312impl From<CliAgentStatus> for AgentStatus {
313    fn from(value: CliAgentStatus) -> Self {
314        match value {
315            CliAgentStatus::Active => Self::Active,
316            CliAgentStatus::Archived => Self::Archived,
317        }
318    }
319}
320
321#[derive(Debug, Subcommand)]
322pub enum ContextCommand {
323    Recent {
324        conversation_id: String,
325        #[arg(long, value_parser = clap::value_parser!(u64).range(1..=1000))]
326        limit: Option<u64>,
327        #[arg(long, value_delimiter = ',')]
328        role: Vec<String>,
329    },
330    Search {
331        conversation_id: String,
332        query: String,
333        #[arg(long)]
334        role: Option<String>,
335        #[arg(long, value_parser = clap::value_parser!(u64).range(1..=1000))]
336        limit: Option<u64>,
337    },
338    /// Append a JSON array of ConversationRecord objects.
339    Append(JsonInput),
340    Maintenance,
341}
342
343#[derive(Debug, Subcommand)]
344pub enum ComputerCommand {
345    /// Create a computer from JSON.
346    Create {
347        #[command(flatten)]
348        input: JsonInput,
349        #[arg(long)]
350        idempotency_key: String,
351    },
352    Get {
353        computer_id: String,
354    },
355    List {
356        #[arg(long)]
357        owner_kind: Option<String>,
358        #[arg(long)]
359        owner_id: Option<String>,
360    },
361    Delete {
362        computer_id: String,
363        #[arg(long)]
364        idempotency_key: String,
365    },
366}
367
368#[derive(Debug, Subcommand)]
369pub enum SpaceCommand {
370    /// Create a space from JSON.
371    Create {
372        #[command(flatten)]
373        input: JsonInput,
374        #[arg(long)]
375        idempotency_key: String,
376    },
377    Get {
378        space_id: String,
379    },
380    List {
381        #[arg(long)]
382        owner_kind: Option<String>,
383        #[arg(long)]
384        owner_id: Option<String>,
385    },
386    Delete {
387        space_id: String,
388        #[arg(long)]
389        idempotency_key: String,
390    },
391}
392
393#[derive(Debug, Subcommand)]
394pub enum ModelCommand {
395    Models,
396    /// Invoke chat/completions with a JSON request document.
397    Chat {
398        #[command(flatten)]
399        input: JsonInput,
400        #[arg(long)]
401        request_id: Option<String>,
402    },
403    /// Invoke responses with a JSON request document.
404    Responses {
405        #[command(flatten)]
406        input: JsonInput,
407        #[arg(long)]
408        request_id: Option<String>,
409    },
410}
411
412#[derive(Debug, Args)]
413pub struct JsonInput {
414    /// JSON file to read; use '-' (the default) for standard input.
415    #[arg(long, short = 'f', default_value = "-")]
416    file: PathBuf,
417}
418
419impl JsonInput {
420    fn read<T: DeserializeOwned>(&self) -> Result<T> {
421        let bytes = read_input(&self.file)?;
422        serde_json::from_slice(&bytes)
423            .with_context(|| format!("invalid JSON input from {}", display_source(&self.file)))
424    }
425}
426
427#[derive(Debug)]
428pub enum CommandOutput {
429    Json(Value),
430    Text(String),
431}
432
433impl CommandOutput {
434    fn json(value: impl Serialize) -> Result<Self> {
435        Ok(Self::Json(serde_json::to_value(value)?))
436    }
437
438    pub fn render(&self, format: OutputFormat) -> Result<String> {
439        match self {
440            Self::Text(value) => Ok(value.clone()),
441            Self::Json(value) => match format {
442                OutputFormat::Pretty => Ok(serde_json::to_string_pretty(value)?),
443                OutputFormat::Json => Ok(serde_json::to_string(value)?),
444            },
445        }
446    }
447}
448
449impl Cli {
450    pub async fn execute(self) -> Result<CommandOutput> {
451        if let Command::Config { command } = &self.command {
452            return execute_config(&self, command);
453        }
454        if let Command::Auth { command } = &self.command {
455            return execute_auth(&self, command);
456        }
457        let client = self.build_client()?;
458        match self.command {
459            Command::Config { .. } | Command::Auth { .. } => {
460                unreachable!("local commands returned before building the client")
461            }
462            Command::Gateway { command } => execute_gateway(&client, command).await,
463            Command::Agent { command } => execute_agent(&client, command).await,
464            Command::Runtime { command } => execute_runtime(&client, command).await,
465            Command::Context { command } => execute_context(&client, command).await,
466            Command::Computer { command } => execute_computer(&client, command).await,
467            Command::Space { command } => execute_space(&client, command).await,
468            Command::Model { command } => execute_model(&client, command).await,
469        }
470    }
471
472    fn build_client(&self) -> Result<InfraClient> {
473        let store = ConfigStore::discover(self.config.as_deref())?;
474        let gateway = store.resolve_gateway(self.gateway.as_deref())?;
475        let (token, _) = resolve_credential(self.token.as_deref(), self.sky_home.as_deref())?;
476        let mut builder = InfraClient::builder(gateway.value)
477            .request_timeout(Duration::from_secs(self.timeout))
478            .trusted_mesh_http(self.trusted_mesh_http)
479            .user_agent(concat!("agent-infra-cli/", env!("CARGO_PKG_VERSION")));
480        if let Some(token) = token {
481            builder = builder.credentials(Arc::new(StaticCredentials::new(token)?));
482        }
483        builder.build().map_err(Into::into)
484    }
485}
486
487fn execute_config(cli: &Cli, command: &ConfigCommand) -> Result<CommandOutput> {
488    let store = ConfigStore::discover(cli.config.as_deref())?;
489    match command {
490        ConfigCommand::Show
491        | ConfigCommand::Gateway {
492            command: ConfigGatewayCommand::Get,
493        } => {
494            let config = store.load()?;
495            let effective = store.resolve_gateway(cli.gateway.as_deref())?;
496            CommandOutput::json(serde_json::json!({
497                "path": store.path(),
498                "gateway": config.gateway,
499                "effectiveGateway": effective.value,
500                "source": effective.source,
501            }))
502        }
503        ConfigCommand::Gateway {
504            command: ConfigGatewayCommand::Set { url },
505        } => {
506            let config = store.set_gateway(url)?;
507            CommandOutput::json(serde_json::json!({
508                "ok": true,
509                "path": store.path(),
510                "gateway": config.gateway,
511            }))
512        }
513        ConfigCommand::Gateway {
514            command: ConfigGatewayCommand::Unset,
515        } => {
516            store.unset_gateway()?;
517            CommandOutput::json(serde_json::json!({
518                "ok": true,
519                "path": store.path(),
520                "gateway": null,
521            }))
522        }
523    }
524}
525
526fn execute_auth(cli: &Cli, command: &AuthCommand) -> Result<CommandOutput> {
527    match command {
528        AuthCommand::Status => {
529            let store = ConfigStore::discover(cli.config.as_deref())?;
530            let gateway = store.resolve_gateway(cli.gateway.as_deref())?;
531            let (_, credential) =
532                resolve_credential(cli.token.as_deref(), cli.sky_home.as_deref())?;
533            CommandOutput::json(serde_json::json!({
534                "gateway": gateway.value,
535                "gatewaySource": gateway.source,
536                "credential": credential,
537            }))
538        }
539    }
540}
541
542async fn execute_gateway(client: &InfraClient, command: GatewayCommand) -> Result<CommandOutput> {
543    match command {
544        GatewayCommand::Discovery => CommandOutput::json(client.gateway().discovery().await?),
545        GatewayCommand::Live => CommandOutput::json(client.gateway().live().await?),
546        GatewayCommand::Ready => CommandOutput::json(client.gateway().ready().await?),
547        GatewayCommand::WaitReady { deadline } => CommandOutput::json(
548            client
549                .gateway()
550                .wait_until_ready(Duration::from_secs(deadline))
551                .await?,
552        ),
553        GatewayCommand::Dependencies => CommandOutput::json(client.gateway().dependencies().await?),
554        GatewayCommand::Routes => CommandOutput::json(client.gateway().routes().await?),
555        GatewayCommand::Metrics => CommandOutput::json(client.gateway().metrics().await?),
556        GatewayCommand::Policy => CommandOutput::json(client.gateway().effective_policy().await?),
557    }
558}
559
560async fn execute_agent(client: &InfraClient, command: AgentCommand) -> Result<CommandOutput> {
561    match command {
562        AgentCommand::Create(input) => {
563            let request: CreateAgentRequest = input.read()?;
564            CommandOutput::json(client.agents().create(&request).await?)
565        }
566        AgentCommand::Get { agent_id } => {
567            let agent_id = parse_agent_id(&agent_id)?;
568            CommandOutput::json(client.agents().get(&agent_id).await?)
569        }
570        AgentCommand::List(args) => {
571            let request = ListAgentsRequest {
572                status: args.status.map(Into::into),
573                related_agent_id: args
574                    .related_agent_id
575                    .as_deref()
576                    .map(parse_agent_id)
577                    .transpose()?,
578                related_runtime_id: args.related_runtime_id,
579                related_user_id: args.related_user_id,
580                cursor: args.cursor,
581                limit: args.limit,
582            };
583            CommandOutput::json(client.agents().list(&request).await?)
584        }
585        AgentCommand::Update { agent_id, input } => {
586            let agent_id = parse_agent_id(&agent_id)?;
587            let request: UpdateAgentRequest = input.read()?;
588            CommandOutput::json(client.agents().update(&agent_id, &request).await?)
589        }
590    }
591}
592
593async fn execute_runtime(client: &InfraClient, command: RuntimeCommand) -> Result<CommandOutput> {
594    match command {
595        RuntimeCommand::Create(args) => {
596            let request = match args.file {
597                Some(file) => JsonInput { file }.read::<CreateRuntimeIdentityRequest>()?,
598                None => CreateRuntimeIdentityRequest {
599                    id: None,
600                    name: args
601                        .name
602                        .ok_or_else(|| anyhow!("--name is required for inline runtime creation"))?,
603                    kind: args.kind.unwrap_or(CliRuntimeKind::Code).into(),
604                    agent_id: None,
605                    relations: Default::default(),
606                    metadata: Default::default(),
607                    idempotency_key: args.idempotency_key.unwrap_or_default(),
608                },
609            };
610            CommandOutput::json(client.runtime_identity().create(&request).await?)
611        }
612        RuntimeCommand::Get { runtime_id } => {
613            let runtime_id = parse_runtime_id(&runtime_id)?;
614            CommandOutput::json(client.runtime_identity().get(&runtime_id).await?)
615        }
616        RuntimeCommand::Rename {
617            runtime_id,
618            name,
619            expected_revision,
620            idempotency_key,
621        } => {
622            let runtime_id = parse_runtime_id(&runtime_id)?;
623            let request = UpdateRuntimeIdentityRequest {
624                expected_revision,
625                name,
626                idempotency_key: idempotency_key.unwrap_or_default(),
627            };
628            CommandOutput::json(
629                client
630                    .runtime_identity()
631                    .update(&runtime_id, &request)
632                    .await?,
633            )
634        }
635        RuntimeCommand::Activate { runtime_id, input } => {
636            execute_runtime_transition(client, runtime_id, "activate", input).await
637        }
638        RuntimeCommand::Suspend { runtime_id, input } => {
639            execute_runtime_transition(client, runtime_id, "suspend", input).await
640        }
641        RuntimeCommand::Resume { runtime_id, input } => {
642            execute_runtime_transition(client, runtime_id, "resume", input).await
643        }
644        RuntimeCommand::Revoke { runtime_id, input } => {
645            execute_runtime_transition(client, runtime_id, "revoke", input).await
646        }
647        RuntimeCommand::Permissions { runtime_id, input } => {
648            let runtime_id = parse_runtime_id(&runtime_id)?;
649            let request: ReplacePermissionsRequest = input.read()?;
650            CommandOutput::json(
651                client
652                    .runtime_identity()
653                    .replace_permissions(&runtime_id, &request)
654                    .await?,
655            )
656        }
657        RuntimeCommand::BootstrapCredential { input } => {
658            let request: IssueRuntimeCredentialRequest = input.read()?;
659            CommandOutput::json(
660                client
661                    .gateway()
662                    .bootstrap_runtime_credential(&request)
663                    .await?,
664            )
665        }
666    }
667}
668
669async fn execute_runtime_transition(
670    client: &InfraClient,
671    runtime_id: String,
672    action: &str,
673    input: JsonInput,
674) -> Result<CommandOutput> {
675    let runtime_id = parse_runtime_id(&runtime_id)?;
676    let request: TransitionIdentityRequest = input.read()?;
677    CommandOutput::json(
678        client
679            .runtime_identity()
680            .transition(&runtime_id, action, &request)
681            .await?,
682    )
683}
684
685async fn execute_context(client: &InfraClient, command: ContextCommand) -> Result<CommandOutput> {
686    match command {
687        ContextCommand::Recent {
688            conversation_id,
689            limit,
690            role,
691        } => {
692            let roles = role.iter().map(String::as_str).collect::<Vec<_>>();
693            CommandOutput::json(
694                client
695                    .context()
696                    .recent_messages(&conversation_id, limit.map(|value| value as usize), &roles)
697                    .await?,
698            )
699        }
700        ContextCommand::Search {
701            conversation_id,
702            query,
703            role,
704            limit,
705        } => CommandOutput::json(
706            client
707                .context()
708                .search_messages(
709                    &conversation_id,
710                    &query,
711                    role.as_deref(),
712                    limit.map(|value| value as usize),
713                )
714                .await?,
715        ),
716        ContextCommand::Append(input) => {
717            let messages: Vec<ConversationRecord> = input.read()?;
718            if messages.is_empty() {
719                bail!("append input must contain at least one message");
720            }
721            CommandOutput::json(client.context().append_messages(messages).await?)
722        }
723        ContextCommand::Maintenance => {
724            CommandOutput::json(client.context().maintenance_status().await?)
725        }
726    }
727}
728
729async fn execute_computer(client: &InfraClient, command: ComputerCommand) -> Result<CommandOutput> {
730    let environment = client.environment();
731    match command {
732        ComputerCommand::Create {
733            input,
734            idempotency_key,
735        } => {
736            let request: CreateComputerRequest = input.read()?;
737            CommandOutput::json(
738                environment
739                    .create_computer(&request, &idempotency_key)
740                    .await?,
741            )
742        }
743        ComputerCommand::Get { computer_id } => {
744            CommandOutput::json(environment.computer(&computer_id).await?)
745        }
746        ComputerCommand::List {
747            owner_kind,
748            owner_id,
749        } => CommandOutput::json(
750            environment
751                .computers(Some(&list_environment_query(owner_kind, owner_id)?))
752                .await?,
753        ),
754        ComputerCommand::Delete {
755            computer_id,
756            idempotency_key,
757        } => CommandOutput::json(
758            environment
759                .delete_computer(&computer_id, &idempotency_key)
760                .await?,
761        ),
762    }
763}
764
765async fn execute_space(client: &InfraClient, command: SpaceCommand) -> Result<CommandOutput> {
766    let environment = client.environment();
767    match command {
768        SpaceCommand::Create {
769            input,
770            idempotency_key,
771        } => {
772            let request: CreateSpaceRequest = input.read()?;
773            CommandOutput::json(
774                environment.create_space(&request, &idempotency_key).await?,
775            )
776        }
777        SpaceCommand::Get { space_id } => {
778            CommandOutput::json(environment.space(&space_id).await?)
779        }
780        SpaceCommand::List {
781            owner_kind,
782            owner_id,
783        } => CommandOutput::json(
784            environment
785                .spaces(Some(&list_environment_query(owner_kind, owner_id)?))
786                .await?,
787        ),
788        SpaceCommand::Delete {
789            space_id,
790            idempotency_key,
791        } => CommandOutput::json(
792            environment
793                .delete_space(&space_id, &idempotency_key)
794                .await?,
795        ),
796    }
797}
798
799fn list_environment_query(
800    owner_kind: Option<String>,
801    owner_id: Option<String>,
802) -> Result<ListEnvironmentResourcesQuery> {
803    let owner_kind = match owner_kind {
804        Some(value) => Some(parse_owner_kind(&value)?),
805        None => None,
806    };
807    if owner_kind.is_some() != owner_id.is_some() {
808        bail!("--owner-kind and --owner-id must be provided together");
809    }
810    Ok(ListEnvironmentResourcesQuery {
811        owner_kind,
812        owner_id,
813        contains_all: vec![],
814    })
815}
816
817fn parse_owner_kind(value: &str) -> Result<OwnerKind> {
818    match value.to_ascii_lowercase().as_str() {
819        "user" => Ok(OwnerKind::User),
820        "agent" => Ok(OwnerKind::Agent),
821        "runtime" => Ok(OwnerKind::Runtime),
822        other => bail!("unsupported owner kind {other:?}; expected user, agent, or runtime"),
823    }
824}
825
826async fn execute_model(client: &InfraClient, command: ModelCommand) -> Result<CommandOutput> {
827    match command {
828        ModelCommand::Models => CommandOutput::json(client.model().models().await?),
829        ModelCommand::Chat { input, request_id } => {
830            let request: Value = input.read()?;
831            ensure_json_object(&request)?;
832            CommandOutput::json(
833                client
834                    .model()
835                    .chat_completions(&request, call_options(request_id))
836                    .await?,
837            )
838        }
839        ModelCommand::Responses { input, request_id } => {
840            let request: Value = input.read()?;
841            ensure_json_object(&request)?;
842            CommandOutput::json(
843                client
844                    .model()
845                    .responses(&request, call_options(request_id))
846                    .await?,
847            )
848        }
849    }
850}
851
852fn call_options(request_id: Option<String>) -> CallOptions {
853    request_id.map_or_else(CallOptions::default, |value| {
854        CallOptions::default().request_id(value)
855    })
856}
857
858fn ensure_json_object(value: &Value) -> Result<()> {
859    if !value.is_object() {
860        bail!("model request must be a JSON object");
861    }
862    Ok(())
863}
864
865fn parse_agent_id(value: &str) -> Result<AgentId> {
866    AgentId::new(value).map_err(|message| anyhow!(message))
867}
868
869fn parse_runtime_id(value: &str) -> Result<RuntimeIdentityId> {
870    RuntimeIdentityId::new(value).map_err(|message| anyhow!(message))
871}
872
873fn read_input(path: &Path) -> Result<Vec<u8>> {
874    if path == Path::new("-") {
875        let mut bytes = Vec::new();
876        io::stdin()
877            .read_to_end(&mut bytes)
878            .context("failed to read standard input")?;
879        Ok(bytes)
880    } else {
881        fs::read(path).with_context(|| format!("failed to read {}", path.display()))
882    }
883}
884
885fn display_source(path: &Path) -> String {
886    if path == Path::new("-") {
887        "standard input".into()
888    } else {
889        path.display().to_string()
890    }
891}
892
893#[cfg(test)]
894mod tests {
895    use super::*;
896    use clap::CommandFactory;
897    use wiremock::{
898        Mock, MockServer, ResponseTemplate,
899        matchers::{body_json, header, method, path},
900    };
901
902    #[test]
903    fn clap_definition_is_valid() {
904        Cli::command().debug_assert();
905    }
906
907    #[test]
908    fn parses_global_options_after_subcommands() {
909        let cli = Cli::try_parse_from([
910            "infra",
911            "gateway",
912            "wait-ready",
913            "--deadline",
914            "42",
915            "--gateway",
916            "https://infra.example.com",
917            "--output",
918            "json",
919        ])
920        .unwrap();
921        assert_eq!(cli.gateway.as_deref(), Some("https://infra.example.com"));
922        assert_eq!(cli.output, OutputFormat::Json);
923        assert!(matches!(
924            cli.command,
925            Command::Gateway {
926                command: GatewayCommand::WaitReady { deadline: 42 }
927            }
928        ));
929    }
930
931    #[test]
932    fn rejects_removed_workspace_subcommand() {
933        let result = Cli::try_parse_from(["infra", "workspace", "write", "README.md"]);
934        assert!(result.is_err());
935    }
936
937    #[test]
938    fn reads_json_request_from_file() {
939        let directory = tempfile::tempdir().unwrap();
940        let path = directory.path().join("request.json");
941        fs::write(&path, br#"{"model":"test"}"#).unwrap();
942        let input = JsonInput { file: path };
943        let value: Value = input.read().unwrap();
944        assert_eq!(value["model"], "test");
945    }
946
947    #[tokio::test]
948    async fn persists_and_resolves_gateway_without_network_access() {
949        let directory = tempfile::tempdir().unwrap();
950        let config = directory.path().join("infra.json");
951        let cli = Cli::try_parse_from([
952            "infra",
953            "config",
954            "gateway",
955            "set",
956            "https://infra.example.com/",
957            "--config",
958            config.to_str().unwrap(),
959        ])
960        .unwrap();
961        let output = cli.execute().await.unwrap();
962        let CommandOutput::Json(value) = output else {
963            panic!("config set must produce JSON");
964        };
965        assert_eq!(value["gateway"], "https://infra.example.com");
966
967        let store = ConfigStore::discover(Some(&config)).unwrap();
968        assert_eq!(
969            store.load().unwrap().gateway.as_deref(),
970            Some("https://infra.example.com")
971        );
972    }
973
974    #[tokio::test]
975    async fn auth_status_detects_encrypted_sky_login_without_exposing_it() {
976        let directory = tempfile::tempdir().unwrap();
977        let sky_home = directory.path().join("sky");
978        fs::create_dir_all(&sky_home).unwrap();
979        fs::write(sky_home.join("credentials"), "ENCv1:secret-ciphertext").unwrap();
980        let config = directory.path().join("infra.json");
981        let cli = Cli::try_parse_from([
982            "infra",
983            "auth",
984            "status",
985            "--config",
986            config.to_str().unwrap(),
987            "--sky-home",
988            sky_home.to_str().unwrap(),
989        ])
990        .unwrap();
991
992        let output = cli.execute().await.unwrap();
993        let rendered = output.render(OutputFormat::Json).unwrap();
994        assert!(rendered.contains("sky_login_encrypted"));
995        assert!(!rendered.contains("secret-ciphertext"));
996    }
997
998    #[test]
999    fn renders_compact_and_pretty_json() {
1000        let output = CommandOutput::Json(serde_json::json!({ "ok": true }));
1001        assert_eq!(output.render(OutputFormat::Json).unwrap(), r#"{"ok":true}"#);
1002        assert!(output.render(OutputFormat::Pretty).unwrap().contains('\n'));
1003    }
1004
1005    #[tokio::test]
1006    async fn model_command_uses_sdk_route_and_bearer_credential() {
1007        let server = MockServer::start().await;
1008        Mock::given(method("GET"))
1009            .and(path("/v1/model/models"))
1010            .and(header("authorization", "Bearer test-token"))
1011            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1012                "object": "list",
1013                "data": [{ "id": "test-model" }]
1014            })))
1015            .expect(1)
1016            .mount(&server)
1017            .await;
1018
1019        let cli = Cli::try_parse_from([
1020            "infra",
1021            "model",
1022            "models",
1023            "--gateway",
1024            &server.uri(),
1025            "--token",
1026            "test-token",
1027        ])
1028        .unwrap();
1029        let output = cli.execute().await.unwrap();
1030        let CommandOutput::Json(value) = output else {
1031            panic!("model list must produce JSON");
1032        };
1033        assert_eq!(value["data"][0]["id"], "test-model");
1034    }
1035
1036    #[tokio::test]
1037    async fn runtime_commands_create_rename_and_revoke_through_the_typed_sdk() {
1038        let server = MockServer::start().await;
1039        let directory = tempfile::tempdir().unwrap();
1040        let transition_request = directory.path().join("transition-runtime.json");
1041        fs::write(
1042            &transition_request,
1043            br#"{"expectedRevision":2,"idempotencyKey":"revoke-runtime-test","reason":"test cleanup"}"#,
1044        )
1045        .unwrap();
1046
1047        Mock::given(method("POST"))
1048            .and(path("/internal/v1/runtime-identities"))
1049            .and(header("authorization", "Bearer test-token"))
1050            .and(body_json(serde_json::json!({
1051                "name": "code-dev",
1052                "kind": "code",
1053                "relations": {},
1054                "metadata": {},
1055                "idempotencyKey": ""
1056            })))
1057            .respond_with(ResponseTemplate::new(201).set_body_json(serde_json::json!({
1058                "identity": runtime_identity_json("code-dev", "pending", 1),
1059                "bootstrapToken": "rti_boot_secret"
1060            })))
1061            .expect(1)
1062            .mount(&server)
1063            .await;
1064        Mock::given(method("PATCH"))
1065            .and(path("/internal/v1/runtime-identities/rti_test"))
1066            .and(header("authorization", "Bearer test-token"))
1067            .and(body_json(serde_json::json!({
1068                "expectedRevision": 1,
1069                "name": "code-dev-renamed",
1070                "idempotencyKey": ""
1071            })))
1072            .respond_with(
1073                ResponseTemplate::new(200).set_body_json(runtime_identity_json(
1074                    "code-dev-renamed",
1075                    "pending",
1076                    2,
1077                )),
1078            )
1079            .expect(1)
1080            .mount(&server)
1081            .await;
1082        Mock::given(method("POST"))
1083            .and(path("/internal/v1/runtime-identities/rti_test:revoke"))
1084            .and(header("authorization", "Bearer test-token"))
1085            .and(body_json(serde_json::json!({
1086                "expectedRevision": 2,
1087                "idempotencyKey": "revoke-runtime-test",
1088                "reason": "test cleanup"
1089            })))
1090            .respond_with(
1091                ResponseTemplate::new(200).set_body_json(runtime_identity_json(
1092                    "code-dev-renamed",
1093                    "revoked",
1094                    3,
1095                )),
1096            )
1097            .expect(1)
1098            .mount(&server)
1099            .await;
1100
1101        let create = Cli::try_parse_from([
1102            "infra",
1103            "runtime",
1104            "create",
1105            "--name",
1106            "code-dev",
1107            "--gateway",
1108            &server.uri(),
1109            "--token",
1110            "test-token",
1111        ])
1112        .unwrap()
1113        .execute()
1114        .await
1115        .unwrap();
1116        let CommandOutput::Json(created) = create else {
1117            panic!("runtime create must produce JSON");
1118        };
1119        assert_eq!(created["identity"]["id"], "rti_test");
1120        assert_eq!(created["identity"]["name"], "code-dev");
1121        assert_eq!(created["bootstrapToken"], "rti_boot_secret");
1122
1123        let rename = Cli::try_parse_from([
1124            "infra",
1125            "runtime",
1126            "rename",
1127            "rti_test",
1128            "code-dev-renamed",
1129            "--expected-revision",
1130            "1",
1131            "--gateway",
1132            &server.uri(),
1133            "--token",
1134            "test-token",
1135        ])
1136        .unwrap()
1137        .execute()
1138        .await
1139        .unwrap();
1140        let CommandOutput::Json(renamed) = rename else {
1141            panic!("runtime rename must produce JSON");
1142        };
1143        assert_eq!(renamed["id"], "rti_test");
1144        assert_eq!(renamed["name"], "code-dev-renamed");
1145        assert_eq!(renamed["revision"], 2);
1146
1147        let revoke = Cli::try_parse_from([
1148            "infra",
1149            "runtime",
1150            "revoke",
1151            "rti_test",
1152            "--file",
1153            transition_request.to_str().unwrap(),
1154            "--gateway",
1155            &server.uri(),
1156            "--token",
1157            "test-token",
1158        ])
1159        .unwrap()
1160        .execute()
1161        .await
1162        .unwrap();
1163        let CommandOutput::Json(revoked) = revoke else {
1164            panic!("runtime revoke must produce JSON");
1165        };
1166        assert_eq!(revoked["status"], "revoked");
1167        assert_eq!(revoked["revision"], 3);
1168    }
1169
1170    #[tokio::test]
1171    async fn runtime_commands_replace_permissions_and_bootstrap_credential() {
1172        let server = MockServer::start().await;
1173        let directory = tempfile::tempdir().unwrap();
1174        let permissions_request = directory.path().join("permissions.json");
1175        fs::write(
1176            &permissions_request,
1177            br#"{"expectedRevision":0,"idempotencyKey":"grant-runtime-test","grants":[{"audience":"agent-infra","capabilities":["gateway:delegation:manage"]}]}"#,
1178        )
1179        .unwrap();
1180        let credential_request = directory.path().join("credential.json");
1181        fs::write(
1182            &credential_request,
1183            br#"{"workloadProof":"bootstrap-proof","audience":"agent-infra","capabilities":["gateway:delegation:manage"]}"#,
1184        )
1185        .unwrap();
1186
1187        Mock::given(method("PUT"))
1188            .and(path(
1189                "/internal/v1/runtime-identities/rti_test/permissions",
1190            ))
1191            .and(header("authorization", "Bearer test-token"))
1192            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1193                "identityId": "rti_test",
1194                "revision": 1,
1195                "grants": [{
1196                    "id": "grt_test",
1197                    "identityId": "rti_test",
1198                    "audience": "agent-infra",
1199                    "capabilities": ["gateway:delegation:manage"],
1200                    "createdAt": "2026-08-13T00:00:00Z"
1201                }]
1202            })))
1203            .expect(1)
1204            .mount(&server)
1205            .await;
1206        Mock::given(method("POST"))
1207            .and(path(
1208                agent_infra_sdk::gateway_contract::RUNTIME_CREDENTIAL_BOOTSTRAP_PATH,
1209            ))
1210            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1211                "accessToken": "runtime-token",
1212                "tokenType": "Bearer",
1213                "expiresInSeconds": 120,
1214                "keyId": "runtime-key"
1215            })))
1216            .expect(1)
1217            .mount(&server)
1218            .await;
1219
1220        let permissions = Cli::try_parse_from([
1221            "infra",
1222            "runtime",
1223            "permissions",
1224            "rti_test",
1225            "--file",
1226            permissions_request.to_str().unwrap(),
1227            "--gateway",
1228            &server.uri(),
1229            "--token",
1230            "test-token",
1231        ])
1232        .unwrap()
1233        .execute()
1234        .await
1235        .unwrap();
1236        let CommandOutput::Json(permissions) = permissions else {
1237            panic!("runtime permissions must produce JSON");
1238        };
1239        assert_eq!(permissions["revision"], 1);
1240
1241        let credential = Cli::try_parse_from([
1242            "infra",
1243            "runtime",
1244            "bootstrap-credential",
1245            "--file",
1246            credential_request.to_str().unwrap(),
1247            "--gateway",
1248            &server.uri(),
1249            "--token",
1250            "test-token",
1251        ])
1252        .unwrap()
1253        .execute()
1254        .await
1255        .unwrap();
1256        let CommandOutput::Json(credential) = credential else {
1257            panic!("runtime bootstrap credential must produce JSON");
1258        };
1259        assert_eq!(credential["accessToken"], "runtime-token");
1260        let requests = server.received_requests().await.unwrap();
1261        let bootstrap = requests
1262            .iter()
1263            .find(|request| {
1264                request.url.path()
1265                    == agent_infra_sdk::gateway_contract::RUNTIME_CREDENTIAL_BOOTSTRAP_PATH
1266            })
1267            .unwrap();
1268        assert!(!bootstrap.headers.contains_key("authorization"));
1269    }
1270
1271    fn runtime_identity_json(name: &str, status: &str, revision: u64) -> Value {
1272        serde_json::json!({
1273            "id": "rti_test",
1274            "name": name,
1275            "scope": { "tenantId": "tenant-test", "projectId": "project-test" },
1276            "kind": "code",
1277            "relations": {
1278                "runtimeIds": ["rti_test"],
1279                "userIds": ["user-test"]
1280            },
1281            "status": status,
1282            "revision": revision,
1283            "grantRevision": 0,
1284            "credentialGeneration": if status == "revoked" { 2 } else { 1 },
1285            "metadata": { "source": "cli-test" },
1286            "createdAt": "2026-08-13T00:00:00Z",
1287            "updatedAt": "2026-08-13T00:00:00Z"
1288        })
1289    }
1290
1291    #[tokio::test]
1292    async fn computer_list_uses_environment_api() {
1293        let server = MockServer::start().await;
1294        Mock::given(method("GET"))
1295            .and(path("/internal/v1/computers"))
1296            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([])))
1297            .expect(1)
1298            .mount(&server)
1299            .await;
1300        let cli = Cli::try_parse_from([
1301            "infra",
1302            "computer",
1303            "list",
1304            "--gateway",
1305            &server.uri(),
1306        ])
1307        .unwrap();
1308
1309        let output = cli.execute().await.unwrap();
1310        let CommandOutput::Json(value) = output else {
1311            panic!("computer list must produce JSON");
1312        };
1313        assert_eq!(value, serde_json::json!([]));
1314    }
1315}
agent_infra_cli/lib.rs

agent_infra_cli/
lib.rs
