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, WaitOptions,
11    agent_contract::{
12        AgentId, AgentStatus, CreateAgentRequest, ListAgentsRequest, UpdateAgentRequest,
13    },
14    context_contract::ConversationRecord,
15    workspace_contract::CreateWorkspaceRequest,
16};
17use agent_workspace_adapter_local::LocalWorkspace;
18use anyhow::{Context, Result, anyhow, bail};
19use clap::{Args, Parser, Subcommand, ValueEnum};
20use serde::{Serialize, de::DeserializeOwned};
21use serde_json::Value;
22
23mod config;
24
25use config::{ConfigStore, resolve_credential};
26
27const DEFAULT_GATEWAY: &str = "http://127.0.0.1:5200";
28
29#[derive(Debug, Parser)]
30#[command(
31    name = "infra",
32    version,
33    about = "Operate Agent Infra through the typed Infra SDK",
34    arg_required_else_help = true
35)]
36pub struct Cli {
37    /// Agent Infra Gateway base URL.
38    #[arg(long, global = true)]
39    pub gateway: Option<String>,
40
41    /// Gateway bearer token. Prefer SKY_GATEWAY_TOKEN over command-line use.
42    #[arg(long, global = true, hide_env_values = true)]
43    pub token: Option<String>,
44
45    /// Override the local Infra CLI config file.
46    #[arg(long, global = true, env = "AGENT_INFRA_CONFIG", value_name = "PATH")]
47    pub config: Option<PathBuf>,
48
49    /// Override the Sky CLI home used for local credential discovery.
50    #[arg(
51        long,
52        global = true,
53        env = "AGENT_INFRA_SKY_HOME",
54        value_name = "PATH",
55        hide = true
56    )]
57    pub sky_home: Option<PathBuf>,
58
59    /// Request timeout in seconds.
60    #[arg(long, global = true, default_value_t = 30, value_parser = clap::value_parser!(u64).range(1..=600))]
61    pub timeout: u64,
62
63    /// Permit plaintext HTTP to non-loopback hosts protected by a trusted service mesh.
64    #[arg(long, global = true)]
65    pub trusted_mesh_http: bool,
66
67    /// JSON rendering mode.
68    #[arg(long, global = true, value_enum, default_value_t = OutputFormat::Pretty)]
69    pub output: OutputFormat,
70
71    /// Bind Workspace file operations directly to this local directory.
72    /// Other Infra services and managed Workspace lifecycle calls still use Gateway.
73    #[arg(
74        long,
75        global = true,
76        env = "AGENT_INFRA_LOCAL_WORKSPACE",
77        value_name = "PATH"
78    )]
79    pub local_workspace: Option<PathBuf>,
80
81    /// Permit `workspace exec` to run host commands inside the local workspace.
82    #[arg(long, global = true, requires = "local_workspace")]
83    pub allow_local_workspace_exec: bool,
84
85    #[command(subcommand)]
86    pub command: Command,
87}
88
89#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, ValueEnum)]
90pub enum OutputFormat {
91    /// Indented JSON.
92    #[default]
93    Pretty,
94    /// Compact JSON.
95    Json,
96}
97
98#[derive(Debug, Subcommand)]
99pub enum Command {
100    /// Persist and inspect local CLI settings.
101    Config {
102        #[command(subcommand)]
103        command: ConfigCommand,
104    },
105    /// Inspect local Gateway credential discovery.
106    Auth {
107        #[command(subcommand)]
108        command: AuthCommand,
109    },
110    /// Inspect Gateway health, discovery, routing, and policy.
111    Gateway {
112        #[command(subcommand)]
113        command: GatewayCommand,
114    },
115    /// Create and inspect Agent Registry resources.
116    Agent {
117        #[command(subcommand)]
118        command: AgentCommand,
119    },
120    /// Read and append durable conversation context.
121    Context {
122        #[command(subcommand)]
123        command: ContextCommand,
124    },
125    /// Operate compatibility files and managed workspaces.
126    Workspace {
127        #[command(subcommand)]
128        command: WorkspaceCommand,
129    },
130    /// Invoke the OpenAI-compatible Model Infra API.
131    Model {
132        #[command(subcommand)]
133        command: ModelCommand,
134    },
135}
136
137#[derive(Debug, Subcommand)]
138pub enum ConfigCommand {
139    /// Show the config file and effective Gateway selection.
140    Show,
141    /// Manage the persisted Gateway address.
142    Gateway {
143        #[command(subcommand)]
144        command: ConfigGatewayCommand,
145    },
146}
147
148#[derive(Debug, Subcommand)]
149pub enum ConfigGatewayCommand {
150    /// Persist a Gateway URL locally.
151    Set { url: String },
152    /// Show the persisted and effective Gateway URL.
153    Get,
154    /// Remove the persisted Gateway URL.
155    Unset,
156}
157
158#[derive(Debug, Subcommand)]
159pub enum AuthCommand {
160    /// Probe explicit, environment, and local Sky credential sources.
161    Status,
162}
163
164#[derive(Debug, Subcommand)]
165pub enum GatewayCommand {
166    Discovery,
167    Live,
168    Ready,
169    /// Poll until the Gateway and its dependencies are ready.
170    WaitReady {
171        #[arg(long, default_value_t = 30, value_parser = clap::value_parser!(u64).range(1..=3600))]
172        deadline: u64,
173    },
174    Dependencies,
175    Routes,
176    Metrics,
177    Policy,
178}
179
180#[derive(Debug, Subcommand)]
181pub enum AgentCommand {
182    /// Create an agent from a CreateAgentRequest JSON document.
183    Create(JsonInput),
184    Get {
185        agent_id: String,
186    },
187    List(AgentListArgs),
188    /// Update an agent from an UpdateAgentRequest JSON document.
189    Update {
190        agent_id: String,
191        #[command(flatten)]
192        input: JsonInput,
193    },
194}
195
196#[derive(Debug, Args)]
197pub struct AgentListArgs {
198    #[arg(long, value_enum)]
199    status: Option<CliAgentStatus>,
200    #[arg(long)]
201    related_agent_id: Option<String>,
202    #[arg(long)]
203    related_runtime_id: Option<String>,
204    #[arg(long)]
205    related_user_id: Option<String>,
206    #[arg(long)]
207    cursor: Option<String>,
208    #[arg(long, value_parser = clap::value_parser!(u32).range(1..=200))]
209    limit: Option<u32>,
210}
211
212#[derive(Clone, Copy, Debug, ValueEnum)]
213enum CliAgentStatus {
214    Active,
215    Archived,
216}
217
218impl From<CliAgentStatus> for AgentStatus {
219    fn from(value: CliAgentStatus) -> Self {
220        match value {
221            CliAgentStatus::Active => Self::Active,
222            CliAgentStatus::Archived => Self::Archived,
223        }
224    }
225}
226
227#[derive(Debug, Subcommand)]
228pub enum ContextCommand {
229    Recent {
230        conversation_id: String,
231        #[arg(long, value_parser = clap::value_parser!(u64).range(1..=1000))]
232        limit: Option<u64>,
233        #[arg(long, value_delimiter = ',')]
234        role: Vec<String>,
235    },
236    Search {
237        conversation_id: String,
238        query: String,
239        #[arg(long)]
240        role: Option<String>,
241        #[arg(long, value_parser = clap::value_parser!(u64).range(1..=1000))]
242        limit: Option<u64>,
243    },
244    /// Append a JSON array of ConversationRecord objects.
245    Append(JsonInput),
246    Maintenance,
247}
248
249#[derive(Debug, Subcommand)]
250pub enum WorkspaceCommand {
251    Capabilities,
252    Read {
253        path: String,
254    },
255    Write(WorkspaceWriteArgs),
256    Exists {
257        path: String,
258    },
259    Mkdir {
260        path: String,
261    },
262    List {
263        path: String,
264    },
265    Exec {
266        command: String,
267        #[arg(long)]
268        cwd: Option<String>,
269    },
270    Remove {
271        path: String,
272    },
273    ListWorkspaces {
274        #[arg(long)]
275        after: Option<String>,
276        #[arg(long, value_parser = clap::value_parser!(u64).range(1..=200))]
277        limit: Option<u64>,
278    },
279    GetWorkspace {
280        workspace_id: String,
281    },
282    WorkspaceUsage {
283        workspace_id: String,
284    },
285    /// Create a managed workspace from JSON.
286    CreateWorkspace {
287        #[command(flatten)]
288        input: JsonInput,
289        #[arg(long)]
290        idempotency_key: String,
291        /// Wait for the asynchronous operation to finish.
292        #[arg(long)]
293        wait: bool,
294        #[arg(long, default_value_t = 900, value_parser = clap::value_parser!(u64).range(1..=604800))]
295        wait_timeout: u64,
296    },
297}
298
299#[derive(Debug, Args)]
300pub struct WorkspaceWriteArgs {
301    pub path: String,
302    /// Literal UTF-8 content. Conflicts with --file.
303    #[arg(long, conflicts_with = "file", required_unless_present = "file")]
304    content: Option<String>,
305    /// Read UTF-8 content from this local file; use '-' for stdin.
306    #[arg(long, conflicts_with = "content", required_unless_present = "content")]
307    file: Option<PathBuf>,
308    /// Perform a conditional write using this revision.
309    #[arg(long)]
310    if_match: Option<String>,
311}
312
313#[derive(Debug, Subcommand)]
314pub enum ModelCommand {
315    Models,
316    /// Invoke chat/completions with a JSON request document.
317    Chat {
318        #[command(flatten)]
319        input: JsonInput,
320        #[arg(long)]
321        request_id: Option<String>,
322    },
323    /// Invoke responses with a JSON request document.
324    Responses {
325        #[command(flatten)]
326        input: JsonInput,
327        #[arg(long)]
328        request_id: Option<String>,
329    },
330}
331
332#[derive(Debug, Args)]
333pub struct JsonInput {
334    /// JSON file to read; use '-' (the default) for standard input.
335    #[arg(long, short = 'f', default_value = "-")]
336    file: PathBuf,
337}
338
339impl JsonInput {
340    fn read<T: DeserializeOwned>(&self) -> Result<T> {
341        let bytes = read_input(&self.file)?;
342        serde_json::from_slice(&bytes)
343            .with_context(|| format!("invalid JSON input from {}", display_source(&self.file)))
344    }
345}
346
347#[derive(Debug)]
348pub enum CommandOutput {
349    Json(Value),
350    Text(String),
351}
352
353impl CommandOutput {
354    fn json(value: impl Serialize) -> Result<Self> {
355        Ok(Self::Json(serde_json::to_value(value)?))
356    }
357
358    pub fn render(&self, format: OutputFormat) -> Result<String> {
359        match self {
360            Self::Text(value) => Ok(value.clone()),
361            Self::Json(value) => match format {
362                OutputFormat::Pretty => Ok(serde_json::to_string_pretty(value)?),
363                OutputFormat::Json => Ok(serde_json::to_string(value)?),
364            },
365        }
366    }
367}
368
369impl Cli {
370    pub async fn execute(self) -> Result<CommandOutput> {
371        if let Command::Config { command } = &self.command {
372            return execute_config(&self, command);
373        }
374        if let Command::Auth { command } = &self.command {
375            return execute_auth(&self, command);
376        }
377        let client = self.build_client()?;
378        match self.command {
379            Command::Config { .. } | Command::Auth { .. } => {
380                unreachable!("local commands returned before building the client")
381            }
382            Command::Gateway { command } => execute_gateway(&client, command).await,
383            Command::Agent { command } => execute_agent(&client, command).await,
384            Command::Context { command } => execute_context(&client, command).await,
385            Command::Workspace { command } => execute_workspace(&client, command).await,
386            Command::Model { command } => execute_model(&client, command).await,
387        }
388    }
389
390    fn build_client(&self) -> Result<InfraClient> {
391        let store = ConfigStore::discover(self.config.as_deref())?;
392        let gateway = store.resolve_gateway(self.gateway.as_deref())?;
393        let (token, _) = resolve_credential(self.token.as_deref(), self.sky_home.as_deref())?;
394        let mut builder = InfraClient::builder(gateway.value)
395            .request_timeout(Duration::from_secs(self.timeout))
396            .trusted_mesh_http(self.trusted_mesh_http)
397            .user_agent(concat!("agent-infra-cli/", env!("CARGO_PKG_VERSION")));
398        if let Some(token) = token {
399            builder = builder.credentials(Arc::new(StaticCredentials::new(token)?));
400        }
401        if let Some(root) = self.local_workspace.as_ref() {
402            let workspace =
403                LocalWorkspace::new_with_host_exec(root.clone(), self.allow_local_workspace_exec);
404            builder = builder.local_workspace(Arc::new(workspace));
405        }
406        builder.build().map_err(Into::into)
407    }
408}
409
410fn execute_config(cli: &Cli, command: &ConfigCommand) -> Result<CommandOutput> {
411    let store = ConfigStore::discover(cli.config.as_deref())?;
412    match command {
413        ConfigCommand::Show
414        | ConfigCommand::Gateway {
415            command: ConfigGatewayCommand::Get,
416        } => {
417            let config = store.load()?;
418            let effective = store.resolve_gateway(cli.gateway.as_deref())?;
419            CommandOutput::json(serde_json::json!({
420                "path": store.path(),
421                "gateway": config.gateway,
422                "effectiveGateway": effective.value,
423                "source": effective.source,
424            }))
425        }
426        ConfigCommand::Gateway {
427            command: ConfigGatewayCommand::Set { url },
428        } => {
429            let config = store.set_gateway(url)?;
430            CommandOutput::json(serde_json::json!({
431                "ok": true,
432                "path": store.path(),
433                "gateway": config.gateway,
434            }))
435        }
436        ConfigCommand::Gateway {
437            command: ConfigGatewayCommand::Unset,
438        } => {
439            store.unset_gateway()?;
440            CommandOutput::json(serde_json::json!({
441                "ok": true,
442                "path": store.path(),
443                "gateway": null,
444            }))
445        }
446    }
447}
448
449fn execute_auth(cli: &Cli, command: &AuthCommand) -> Result<CommandOutput> {
450    match command {
451        AuthCommand::Status => {
452            let store = ConfigStore::discover(cli.config.as_deref())?;
453            let gateway = store.resolve_gateway(cli.gateway.as_deref())?;
454            let (_, credential) =
455                resolve_credential(cli.token.as_deref(), cli.sky_home.as_deref())?;
456            CommandOutput::json(serde_json::json!({
457                "gateway": gateway.value,
458                "gatewaySource": gateway.source,
459                "credential": credential,
460            }))
461        }
462    }
463}
464
465async fn execute_gateway(client: &InfraClient, command: GatewayCommand) -> Result<CommandOutput> {
466    match command {
467        GatewayCommand::Discovery => CommandOutput::json(client.gateway().discovery().await?),
468        GatewayCommand::Live => CommandOutput::json(client.gateway().live().await?),
469        GatewayCommand::Ready => CommandOutput::json(client.gateway().ready().await?),
470        GatewayCommand::WaitReady { deadline } => CommandOutput::json(
471            client
472                .gateway()
473                .wait_until_ready(Duration::from_secs(deadline))
474                .await?,
475        ),
476        GatewayCommand::Dependencies => CommandOutput::json(client.gateway().dependencies().await?),
477        GatewayCommand::Routes => CommandOutput::json(client.gateway().routes().await?),
478        GatewayCommand::Metrics => CommandOutput::json(client.gateway().metrics().await?),
479        GatewayCommand::Policy => CommandOutput::json(client.gateway().effective_policy().await?),
480    }
481}
482
483async fn execute_agent(client: &InfraClient, command: AgentCommand) -> Result<CommandOutput> {
484    match command {
485        AgentCommand::Create(input) => {
486            let request: CreateAgentRequest = input.read()?;
487            CommandOutput::json(client.agents().create(&request).await?)
488        }
489        AgentCommand::Get { agent_id } => {
490            let agent_id = parse_agent_id(&agent_id)?;
491            CommandOutput::json(client.agents().get(&agent_id).await?)
492        }
493        AgentCommand::List(args) => {
494            let request = ListAgentsRequest {
495                status: args.status.map(Into::into),
496                related_agent_id: args
497                    .related_agent_id
498                    .as_deref()
499                    .map(parse_agent_id)
500                    .transpose()?,
501                related_runtime_id: args.related_runtime_id,
502                related_user_id: args.related_user_id,
503                cursor: args.cursor,
504                limit: args.limit,
505            };
506            CommandOutput::json(client.agents().list(&request).await?)
507        }
508        AgentCommand::Update { agent_id, input } => {
509            let agent_id = parse_agent_id(&agent_id)?;
510            let request: UpdateAgentRequest = input.read()?;
511            CommandOutput::json(client.agents().update(&agent_id, &request).await?)
512        }
513    }
514}
515
516async fn execute_context(client: &InfraClient, command: ContextCommand) -> Result<CommandOutput> {
517    match command {
518        ContextCommand::Recent {
519            conversation_id,
520            limit,
521            role,
522        } => {
523            let roles = role.iter().map(String::as_str).collect::<Vec<_>>();
524            CommandOutput::json(
525                client
526                    .context()
527                    .recent_messages(&conversation_id, limit.map(|value| value as usize), &roles)
528                    .await?,
529            )
530        }
531        ContextCommand::Search {
532            conversation_id,
533            query,
534            role,
535            limit,
536        } => CommandOutput::json(
537            client
538                .context()
539                .search_messages(
540                    &conversation_id,
541                    &query,
542                    role.as_deref(),
543                    limit.map(|value| value as usize),
544                )
545                .await?,
546        ),
547        ContextCommand::Append(input) => {
548            let messages: Vec<ConversationRecord> = input.read()?;
549            if messages.is_empty() {
550                bail!("append input must contain at least one message");
551            }
552            CommandOutput::json(client.context().append_messages(messages).await?)
553        }
554        ContextCommand::Maintenance => {
555            CommandOutput::json(client.context().maintenance_status().await?)
556        }
557    }
558}
559
560async fn execute_workspace(
561    client: &InfraClient,
562    command: WorkspaceCommand,
563) -> Result<CommandOutput> {
564    let workspace = client.workspace();
565    match command {
566        WorkspaceCommand::Capabilities => CommandOutput::json(workspace.capabilities().await?),
567        WorkspaceCommand::Read { path } => {
568            Ok(CommandOutput::Text(workspace.read_file(&path).await?))
569        }
570        WorkspaceCommand::Write(args) => {
571            let content = match (args.content, args.file) {
572                (Some(content), None) => content,
573                (None, Some(path)) => String::from_utf8(read_input(&path)?)
574                    .with_context(|| format!("{} is not valid UTF-8", display_source(&path)))?,
575                _ => return Err(anyhow!("exactly one of --content or --file is required")),
576            };
577            let revision = workspace
578                .write_file_with_options(
579                    &args.path,
580                    &content,
581                    args.if_match.as_deref(),
582                    CallOptions::default(),
583                )
584                .await?;
585            CommandOutput::json(serde_json::json!({ "ok": true, "revision": revision }))
586        }
587        WorkspaceCommand::Exists { path } => {
588            CommandOutput::json(serde_json::json!({ "exists": workspace.exists(&path).await? }))
589        }
590        WorkspaceCommand::Mkdir { path } => {
591            workspace.create_dir_all(&path).await?;
592            CommandOutput::json(serde_json::json!({ "ok": true }))
593        }
594        WorkspaceCommand::List { path } => CommandOutput::json(workspace.list_dir(&path).await?),
595        WorkspaceCommand::Exec { command, cwd } => {
596            CommandOutput::json(workspace.exec(&command, cwd.as_deref()).await?)
597        }
598        WorkspaceCommand::Remove { path } => {
599            workspace.remove_file(&path).await?;
600            CommandOutput::json(serde_json::json!({ "ok": true }))
601        }
602        WorkspaceCommand::ListWorkspaces { after, limit } => CommandOutput::json(
603            workspace
604                .list_workspaces(after.as_deref(), limit.map(|value| value as usize))
605                .await?,
606        ),
607        WorkspaceCommand::GetWorkspace { workspace_id } => {
608            CommandOutput::json(workspace.workspace(&workspace_id).await?)
609        }
610        WorkspaceCommand::WorkspaceUsage { workspace_id } => {
611            CommandOutput::json(workspace.workspace_usage(&workspace_id).await?)
612        }
613        WorkspaceCommand::CreateWorkspace {
614            input,
615            idempotency_key,
616            wait,
617            wait_timeout,
618        } => {
619            let request: CreateWorkspaceRequest = input.read()?;
620            let handle = workspace
621                .create_workspace(&request, &idempotency_key)
622                .await?;
623            if wait {
624                CommandOutput::json(
625                    handle
626                        .wait(WaitOptions::deadline(Duration::from_secs(wait_timeout)))
627                        .await?,
628                )
629            } else {
630                CommandOutput::json(serde_json::json!({ "operationId": handle.id() }))
631            }
632        }
633    }
634}
635
636async fn execute_model(client: &InfraClient, command: ModelCommand) -> Result<CommandOutput> {
637    match command {
638        ModelCommand::Models => CommandOutput::json(client.model().models().await?),
639        ModelCommand::Chat { input, request_id } => {
640            let request: Value = input.read()?;
641            ensure_json_object(&request)?;
642            CommandOutput::json(
643                client
644                    .model()
645                    .chat_completions(&request, call_options(request_id))
646                    .await?,
647            )
648        }
649        ModelCommand::Responses { input, request_id } => {
650            let request: Value = input.read()?;
651            ensure_json_object(&request)?;
652            CommandOutput::json(
653                client
654                    .model()
655                    .responses(&request, call_options(request_id))
656                    .await?,
657            )
658        }
659    }
660}
661
662fn call_options(request_id: Option<String>) -> CallOptions {
663    request_id.map_or_else(CallOptions::default, |value| {
664        CallOptions::default().request_id(value)
665    })
666}
667
668fn ensure_json_object(value: &Value) -> Result<()> {
669    if !value.is_object() {
670        bail!("model request must be a JSON object");
671    }
672    Ok(())
673}
674
675fn parse_agent_id(value: &str) -> Result<AgentId> {
676    AgentId::new(value).map_err(|message| anyhow!(message))
677}
678
679fn read_input(path: &Path) -> Result<Vec<u8>> {
680    if path == Path::new("-") {
681        let mut bytes = Vec::new();
682        io::stdin()
683            .read_to_end(&mut bytes)
684            .context("failed to read standard input")?;
685        Ok(bytes)
686    } else {
687        fs::read(path).with_context(|| format!("failed to read {}", path.display()))
688    }
689}
690
691fn display_source(path: &Path) -> String {
692    if path == Path::new("-") {
693        "standard input".into()
694    } else {
695        path.display().to_string()
696    }
697}
698
699#[cfg(test)]
700mod tests {
701    use super::*;
702    use clap::CommandFactory;
703    use wiremock::{
704        Mock, MockServer, ResponseTemplate,
705        matchers::{header, method, path},
706    };
707
708    #[test]
709    fn clap_definition_is_valid() {
710        Cli::command().debug_assert();
711    }
712
713    #[test]
714    fn parses_global_options_after_subcommands() {
715        let cli = Cli::try_parse_from([
716            "infra",
717            "gateway",
718            "wait-ready",
719            "--deadline",
720            "42",
721            "--gateway",
722            "https://infra.example.com",
723            "--output",
724            "json",
725        ])
726        .unwrap();
727        assert_eq!(cli.gateway.as_deref(), Some("https://infra.example.com"));
728        assert_eq!(cli.output, OutputFormat::Json);
729        assert!(matches!(
730            cli.command,
731            Command::Gateway {
732                command: GatewayCommand::WaitReady { deadline: 42 }
733            }
734        ));
735    }
736
737    #[test]
738    fn rejects_ambiguous_workspace_write_content() {
739        let result = Cli::try_parse_from([
740            "infra",
741            "workspace",
742            "write",
743            "README.md",
744            "--content",
745            "one",
746            "--file",
747            "two.txt",
748        ]);
749        assert!(result.is_err());
750    }
751
752    #[test]
753    fn reads_json_request_from_file() {
754        let directory = tempfile::tempdir().unwrap();
755        let path = directory.path().join("request.json");
756        fs::write(&path, br#"{"model":"test"}"#).unwrap();
757        let input = JsonInput { file: path };
758        let value: Value = input.read().unwrap();
759        assert_eq!(value["model"], "test");
760    }
761
762    #[tokio::test]
763    async fn persists_and_resolves_gateway_without_network_access() {
764        let directory = tempfile::tempdir().unwrap();
765        let config = directory.path().join("infra.json");
766        let cli = Cli::try_parse_from([
767            "infra",
768            "config",
769            "gateway",
770            "set",
771            "https://infra.example.com/",
772            "--config",
773            config.to_str().unwrap(),
774        ])
775        .unwrap();
776        let output = cli.execute().await.unwrap();
777        let CommandOutput::Json(value) = output else {
778            panic!("config set must produce JSON");
779        };
780        assert_eq!(value["gateway"], "https://infra.example.com");
781
782        let store = ConfigStore::discover(Some(&config)).unwrap();
783        assert_eq!(
784            store.load().unwrap().gateway.as_deref(),
785            Some("https://infra.example.com")
786        );
787    }
788
789    #[tokio::test]
790    async fn auth_status_detects_encrypted_sky_login_without_exposing_it() {
791        let directory = tempfile::tempdir().unwrap();
792        let sky_home = directory.path().join("sky");
793        fs::create_dir_all(&sky_home).unwrap();
794        fs::write(sky_home.join("credentials"), "ENCv1:secret-ciphertext").unwrap();
795        let config = directory.path().join("infra.json");
796        let cli = Cli::try_parse_from([
797            "infra",
798            "auth",
799            "status",
800            "--config",
801            config.to_str().unwrap(),
802            "--sky-home",
803            sky_home.to_str().unwrap(),
804        ])
805        .unwrap();
806
807        let output = cli.execute().await.unwrap();
808        let rendered = output.render(OutputFormat::Json).unwrap();
809        assert!(rendered.contains("sky_login_encrypted"));
810        assert!(!rendered.contains("secret-ciphertext"));
811    }
812
813    #[test]
814    fn renders_compact_and_pretty_json() {
815        let output = CommandOutput::Json(serde_json::json!({ "ok": true }));
816        assert_eq!(output.render(OutputFormat::Json).unwrap(), r#"{"ok":true}"#);
817        assert!(output.render(OutputFormat::Pretty).unwrap().contains('\n'));
818    }
819
820    #[tokio::test]
821    async fn model_command_uses_sdk_route_and_bearer_credential() {
822        let server = MockServer::start().await;
823        Mock::given(method("GET"))
824            .and(path("/v1/model/models"))
825            .and(header("authorization", "Bearer test-token"))
826            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
827                "object": "list",
828                "data": [{ "id": "test-model" }]
829            })))
830            .expect(1)
831            .mount(&server)
832            .await;
833
834        let cli = Cli::try_parse_from([
835            "infra",
836            "model",
837            "models",
838            "--gateway",
839            &server.uri(),
840            "--token",
841            "test-token",
842        ])
843        .unwrap();
844        let output = cli.execute().await.unwrap();
845        let CommandOutput::Json(value) = output else {
846            panic!("model list must produce JSON");
847        };
848        assert_eq!(value["data"][0]["id"], "test-model");
849    }
850
851    #[tokio::test]
852    async fn local_workspace_data_plane_does_not_call_gateway() {
853        let directory = tempfile::tempdir().unwrap();
854        fs::write(directory.path().join("input.txt"), "local-only").unwrap();
855
856        let cli = Cli::try_parse_from([
857            "infra",
858            "workspace",
859            "read",
860            "input.txt",
861            "--local-workspace",
862            directory.path().to_str().unwrap(),
863            "--gateway",
864            "https://gateway.invalid",
865        ])
866        .unwrap();
867        let client = cli.build_client().unwrap();
868
869        assert_eq!(
870            client.workspace().read_file("input.txt").await.unwrap(),
871            "local-only"
872        );
873        client.workspace().create_dir_all("nested").await.unwrap();
874        let revision = client
875            .workspace()
876            .write_file_with_options(
877                "nested/output.txt",
878                "written-directly",
879                None,
880                CallOptions::default(),
881            )
882            .await
883            .unwrap();
884        assert!(revision.is_some());
885        assert!(
886            client
887                .workspace()
888                .exists("nested/output.txt")
889                .await
890                .unwrap()
891        );
892        assert_eq!(
893            fs::read_to_string(directory.path().join("nested/output.txt")).unwrap(),
894            "written-directly"
895        );
896        assert!(client.workspace().is_dir("nested").await.unwrap());
897        assert_eq!(
898            client
899                .workspace()
900                .read_file_bytes("nested/output.txt")
901                .await
902                .unwrap(),
903            b"written-directly"
904        );
905        assert!(
906            client
907                .workspace()
908                .walk_tree(".", 4)
909                .await
910                .unwrap()
911                .iter()
912                .any(|path| path.ends_with("nested/output.txt"))
913        );
914        client
915            .workspace()
916            .remove_file_if_match("nested/output.txt", revision.as_deref().unwrap())
917            .await
918            .unwrap();
919        assert!(
920            !client
921                .workspace()
922                .exists("nested/output.txt")
923                .await
924                .unwrap()
925        );
926    }
927
928    #[tokio::test]
929    async fn local_workspace_binding_keeps_managed_lifecycle_on_gateway() {
930        let server = MockServer::start().await;
931        Mock::given(method("GET"))
932            .and(path("/internal/v1/workspaces"))
933            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
934                "items": [],
935                "nextCursor": null
936            })))
937            .expect(1)
938            .mount(&server)
939            .await;
940        let directory = tempfile::tempdir().unwrap();
941        let cli = Cli::try_parse_from([
942            "infra",
943            "workspace",
944            "list-workspaces",
945            "--gateway",
946            &server.uri(),
947            "--local-workspace",
948            directory.path().to_str().unwrap(),
949        ])
950        .unwrap();
951
952        let output = cli.execute().await.unwrap();
953        let CommandOutput::Json(value) = output else {
954            panic!("managed workspace list must produce JSON");
955        };
956        assert_eq!(value["items"], serde_json::json!([]));
957    }
958}
agent_infra_cli/lib.rs

agent_infra_cli/
lib.rs
