From 7fbca0ada6f9efe9bf66fc8ba4078f23b8fd7229 Mon Sep 17 00:00:00 2001 From: okaybro <66475772+Chaitanya-Keyal@users.noreply.github.com> Date: Tue, 12 May 2026 03:09:42 +0530 Subject: [PATCH] feat: add `--with-status` flag to `load` --- README.md | 4 +- src/cli.rs | 3 ++ src/lib.rs | 7 ++- src/profiles.rs | 117 ++++++++++++++++++++++++++++++++++-------------- tests/cli.rs | 60 +++++++++++++++++++++++++ 5 files changed, 155 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 9429282..5cc2017 100644 --- a/README.md +++ b/README.md @@ -93,8 +93,8 @@ codex-profiles load --label work --force Save current auth.json
Optional label - codex-profiles load
(--label <name> | --id <profile-id>)
[--force] - Load a saved profile
Choose a target profile and force when needed + codex-profiles load
(--label <name> | --id <profile-id>)
[--force] [--with-status] + Load a saved profile
Optionally show status after loading codex-profiles list
[--show-id] [--json] diff --git a/src/cli.rs b/src/cli.rs index 2e0cf78..a1ca74b 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -38,6 +38,9 @@ pub enum Commands { /// Continue without saving the current unsaved profile first #[arg(long)] force: bool, + /// Show status for the loaded profile after a successful load + #[arg(long)] + with_status: bool, }, /// List saved profiles List { diff --git a/src/lib.rs b/src/lib.rs index d443ecc..abe30e1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -72,7 +72,12 @@ fn run(cli: Cli) -> Result<(), String> { match cli.command { Commands::Save { label } => save_profile(&paths, label, json), - Commands::Load { label, id, force } => load_profile(&paths, label, id, force, json), + Commands::Load { + label, + id, + force, + with_status, + } => load_profile(&paths, label, id, force, with_status, json), Commands::List { show_id } => list_profiles(&paths, json, show_id), Commands::Export { label, id, output } => export_profiles(&paths, label, id, output, json), Commands::Import { input } => import_profiles(&paths, input, json), diff --git a/src/profiles.rs b/src/profiles.rs index 04223d8..a1efc46 100644 --- a/src/profiles.rs +++ b/src/profiles.rs @@ -389,6 +389,7 @@ pub fn load_profile( label: Option, id: Option, force: bool, + with_status: bool, json: bool, ) -> Result<(), String> { let use_color_err = use_color_stderr(); @@ -464,15 +465,25 @@ pub fn load_profile( label.clone(), ); store.save(paths)?; + drop(store); + + let mut profile_json = serde_json::json!({ + "id": selected_id, + "label": label, + }); + if json && with_status { + match current_status_json_value(paths) { + Ok(status) => { + profile_json["status"] = status; + } + Err(err) => { + profile_json["status_error"] = serde_json::Value::String(normalize_error(&err)); + } + } + } if json { - let result = CommandResultJson::success( - "load", - serde_json::json!({ - "id": selected_id, - "label": label, - }), - ); + let result = CommandResultJson::success("load", profile_json); result.print()?; return Ok(()); } @@ -482,6 +493,13 @@ pub fn load_profile( use_color_out, ); print_output_block(&message); + if with_status && let Err(err) = loaded_profile_status(paths) { + let message = format!( + "Loaded profile, but status retrieval failed: {}", + normalize_error(&err) + ); + eprintln!("{}", format_warning(&message, use_color_err)); + } Ok(()) } @@ -634,6 +652,25 @@ pub fn status_profiles( return status_selected_profile(paths, label.as_deref(), id.as_deref(), json); } + status_current_profile(paths, json) +} + +fn status_current_profile(paths: &Paths, json: bool) -> Result<(), String> { + let (current_entry, ctx) = current_status_entry(paths, json)?; + if json { + return print_current_status_json(current_entry); + } + if let Some(entry) = current_entry { + let lines = render_entries(std::slice::from_ref(&entry), &ctx, false); + print_output_block(&lines.join("\n")); + } else { + let message = format_no_profiles(paths, ctx.use_color); + print_output_block(&message); + } + Ok(()) +} + +fn current_status_entry(paths: &Paths, json: bool) -> Result<(Option, ListCtx), String> { let snapshot = load_snapshot(paths, false)?; let current_saved_id = current_saved_id(paths, &snapshot.tokens); let mut ctx = ListCtx::new(paths, true, false, false); @@ -642,17 +679,28 @@ pub fn status_profiles( } let labels = &snapshot.labels; let tokens_map = &snapshot.tokens; - let current_entry = make_current(paths, current_saved_id.as_deref(), labels, tokens_map, &ctx); - if json { - return print_current_status_json(current_entry); - } - if let Some(entry) = current_entry { - let lines = render_entries(&[entry], &ctx, false); - print_output_block(&lines.join("\n")); - } else { + let current = make_current(paths, current_saved_id.as_deref(), labels, tokens_map, &ctx); + Ok((current, ctx)) +} + +fn current_status_json_value(paths: &Paths) -> Result { + let (current, _) = current_status_entry(paths, true)?; + let payload = current.map(status_profile_json); + serde_json::to_value(payload).map_err(|err| crate::msg1(PROFILE_ERR_SERIALIZE_INDEX, err)) +} + +fn loaded_profile_status(paths: &Paths) -> Result<(), String> { + let (current_entry, ctx) = current_status_entry(paths, false)?; + let Some(entry) = current_entry else { let message = format_no_profiles(paths, ctx.use_color); print_output_block(&message); + return Ok(()); + }; + if let Some(err) = entry.error_summary.as_deref() { + return Err(err.to_string()); } + let lines = render_entry_details(&entry); + print_output_block(&lines.join("\n")); Ok(()) } @@ -2001,24 +2049,7 @@ fn render_entries(entries: &[Entry], ctx: &ListCtx, allow_plain_spacing: bool) - } else { entry_lines.push(header); entry_lines.push(String::new()); - entry_lines.extend(entry.details.iter().flat_map(|line| { - if line.is_empty() { - vec![String::new()] - } else { - line.lines() - .enumerate() - .map(|(index, part)| { - if part.is_empty() { - String::new() - } else if index == 0 { - format!(" {part}") - } else { - part.to_string() - } - }) - .collect::>() - } - })); + entry_lines.extend(render_entry_details(entry)); } lines.extend(entry_lines); if idx + 1 < entries.len() { @@ -2031,6 +2062,26 @@ fn render_entries(entries: &[Entry], ctx: &ListCtx, allow_plain_spacing: bool) - lines } +fn render_entry_details(entry: &Entry) -> Vec { + let mut details = Vec::new(); + for line in &entry.details { + if line.is_empty() { + details.push(String::new()); + continue; + } + details.extend(line.lines().enumerate().map(|(index, part)| { + if part.is_empty() { + String::new() + } else if index == 0 { + format!(" {part}") + } else { + part.to_string() + } + })); + } + details +} + fn push_separator(lines: &mut Vec, allow_plain_spacing: bool) { if !is_plain() || allow_plain_spacing { lines.push(String::new()); diff --git a/tests/cli.rs b/tests/cli.rs index 6f69bc1..a050918 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -1400,6 +1400,44 @@ fn ui_load_by_id_command() { assert!(env.read_auth().contains(BETA_ACCOUNT)); } +#[test] +fn ui_load_with_status_command() { + let env = TestEnv::new(); + seed_profiles(&env); + seed_alpha(&env); + let usage_body = r#"{"rate_limit":{"primary_window":{"used_percent":20,"limit_window_seconds":18000,"reset_at":2000000000}}}"#; + let (usage_addr, usage_handle) = start_usage_server(usage_body, 1).expect("usage server"); + env.write_config(&format!("http://{usage_addr}/backend-api")); + + let output = env.run(&["load", "--label", "beta", "--with-status"]); + assert!(output.contains("Loaded profile")); + assert!(output.contains(BETA_EMAIL)); + assert_eq!(output.matches(BETA_EMAIL).count(), 1, "{output}"); + assert!(!output.contains(ALPHA_EMAIL)); + assert!(output.contains("80% left"), "{output}"); + assert!(env.read_auth().contains(BETA_ACCOUNT)); + + let _ = usage_handle.join(); +} + +#[test] +fn ui_load_with_status_keeps_load_success_when_usage_fails() { + let env = TestEnv::new(); + seed_profiles(&env); + seed_alpha(&env); + env.write_config("http://127.0.0.1:1/backend-api"); + + let output = env.run_output(&["load", "--label", "beta", "--with-status"]); + assert!(output.status.success()); + let stdout = ascii_only(String::from_utf8_lossy(&output.stdout).as_ref()); + let stderr = ascii_only(String::from_utf8_lossy(&output.stderr).as_ref()); + assert!(stdout.contains("Loaded profile")); + assert!(stdout.contains(BETA_EMAIL)); + assert!(stderr.contains("Loaded profile, but status retrieval failed")); + assert!(stderr.contains("Could not reach usage service")); + assert!(env.read_auth().contains(BETA_ACCOUNT)); +} + #[test] fn ui_load_by_id_force_skips_unsaved_prompt() { let env = TestEnv::new(); @@ -2780,6 +2818,28 @@ fn json_load_returns_success_shape() { assert!(profile.get("default").is_none()); } +#[test] +fn json_load_with_status_returns_success_shape() { + let env = TestEnv::new(); + seed_alpha(&env); + env.run(&["save", "--label", "alpha"]); + let usage_body = r#"{"rate_limit":{"primary_window":{"used_percent":20,"limit_window_seconds":18000,"reset_at":2000000000}}}"#; + let (usage_addr, usage_handle) = start_usage_server(usage_body, 1).expect("usage server"); + env.write_config(&format!("http://{usage_addr}/backend-api")); + + let raw = env.run(&["load", "--label", "alpha", "--with-status", "--json"]); + let v = parse_json(&raw); + + assert_eq!(v["command"], "load"); + assert_eq!(v["success"], true); + let profile = &v["profile"]; + assert_eq!(profile["label"], "alpha"); + assert_eq!(profile["status"]["id"], ALPHA_ID); + assert_eq!(profile["status"]["usage"]["state"], "ok"); + + let _ = usage_handle.join(); +} + #[test] fn json_delete_returns_success_shape() { let env = TestEnv::new();