diff --git a/Cargo.toml b/Cargo.toml index dc88fd09..85ab163b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -73,7 +73,7 @@ browser = [ [dependencies] # CLI (optional — not needed for browser WASM library) -clap = { version = "4", features = ["derive"], optional = true } +clap = { version = "4", features = ["derive", "string"], optional = true } clap_complete = { version = "4", optional = true } # Async runtime (features selected by native/wasi feature flags) diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index d0dd44ff..ce91a8ec 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -81,6 +81,7 @@ pup [options] # Nested commands | fleet | agents (list, get, versions, tracers), deployments (list, get, configure, upgrade, cancel), schedules (list, get, create, update, delete, trigger), tracers (list), clusters (list), instrumented-pods (list) | src/commands/fleet.rs | ✅ | | skills | list, install, path (positional ``: claude/cursor/codex/opencode/windsurf/gemini/pi/devin/all; `--name`, `--type`, `--project` for project-local scope; bundled skills include supplementary references) | src/commands/skills.rs | ✅ | | runbooks | list, describe, run, import, validate | src/commands/runbooks.rs | ✅ | +| setup | (all arguments passed to AI Setup, run from npm with npx; needs Node.js 22+) | src/commands/setup.rs | ✅ | | workflows | get, create, update, diff, delete, run, instances (list, get, cancel), connections (get, create, update, delete) | src/commands/workflows.rs | ✅ | | investigations | list, get, trigger | src/commands/investigations.rs | ✅ | | change-requests | create, get, update, create-branch, decisions (update, delete) | src/commands/change_management.rs | ✅ | diff --git a/src/commands/agent.rs b/src/commands/agent.rs index 895830c1..74107c22 100644 --- a/src/commands/agent.rs +++ b/src/commands/agent.rs @@ -77,10 +77,23 @@ pub fn returns_for(full_path: &str) -> serde_json::Value { "traces search" => traces_search_returns(), "metrics query" | "metrics search" => metrics_query_returns(), "security findings schema" => findings_schema_returns(), + "setup" => setup_returns(), _ => generic_returns(), } } +fn setup_returns() -> serde_json::Value { + serde_json::json!({ + "documented": false, + "data": {"description": "No envelope and no API body: AI Setup's output is passed through to stdout and stderr."}, + "notes": [ + "Headless runs (an AI agent is detected, or there is no terminal) print AI_SETUP_PROGRESS lines and end with one AI_SETUP_RESULT line: `status=success product= code=0 url=` or `status=error product= code= reason=\"\"`.", + "The exit code is AI Setup's: 0 on success, non-zero on failure.", + "Needs Node.js 22+ and npx; without them pup prints how to install them and exits 1." + ] + }) +} + fn generic_returns() -> serde_json::Value { serde_json::json!({ "documented": false, diff --git a/src/commands/alias.rs b/src/commands/alias.rs index 61c47318..d202905c 100644 --- a/src/commands/alias.rs +++ b/src/commands/alias.rs @@ -78,7 +78,11 @@ pub fn delete(names: Vec) -> Result<()> { /// Core expansion logic, operating on an already-loaded alias map. /// Exposed separately so unit tests can exercise it without touching the /// filesystem. -fn apply_expansion(args: &[String], name: &str, aliases: &BTreeMap) -> Vec { +pub(crate) fn apply_expansion( + args: &[String], + name: &str, + aliases: &BTreeMap, +) -> Vec { // Built-in commands cannot be shadowed by aliases. #[cfg(not(target_arch = "wasm32"))] if crate::extensions::is_builtin_command(name) { diff --git a/src/commands/auth.rs b/src/commands/auth.rs index c0be75cf..190de07a 100644 --- a/src/commands/auth.rs +++ b/src/commands/auth.rs @@ -47,7 +47,7 @@ pub async fn login( scopes: Vec, callback_port: Option, org_uuid: Option<&str>, -) -> Result<()> { +) -> Result { use crate::auth::{dcr, pkce}; let site = &cfg.site; @@ -238,11 +238,12 @@ pub async fn login( let saved_org_uuid = validated_dd_oid .map(String::from) .or_else(|| org_uuid.map(String::from)); - storage::save_session(&storage::SessionEntry { + let session = storage::SessionEntry { site: effective_site.to_string(), org: saved_org.map(String::from), org_uuid: saved_org_uuid, - })?; + }; + storage::save_session(&session)?; // Enforce the single-slot invariant: a bare login (no org) replaces any // previous no-org session on a different site so that @@ -259,7 +260,7 @@ pub async fn login( eprintln!(" Access token expires: {expires_at}"); eprintln!(" Token stored in: {location}"); - Ok(()) + Ok(session) } /// Pick the site to key the exchanged token, session, and DCR credential under. diff --git a/src/commands/mod.rs b/src/commands/mod.rs index b3779ad5..8e5914ca 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -93,6 +93,10 @@ pub mod seats; pub mod security; pub mod service_catalog; #[cfg(not(target_arch = "wasm32"))] +pub mod setup; +#[cfg(not(target_arch = "wasm32"))] +pub mod setup_help; +#[cfg(not(target_arch = "wasm32"))] pub mod skills; #[cfg(not(target_arch = "wasm32"))] pub mod skills_remote; diff --git a/src/commands/setup.rs b/src/commands/setup.rs new file mode 100644 index 00000000..f1b86a8b --- /dev/null +++ b/src/commands/setup.rs @@ -0,0 +1,1152 @@ +use anyhow::Result; +use std::io::IsTerminal; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::time::Duration; + +use crate::auth::storage::SessionEntry; +use crate::commands::setup_help; +use crate::config::Config; +use crate::extensions::exec; + +const AI_SETUP_PACKAGE: &str = "@datadog/ai-setup-cli"; +const AI_SETUP_BIN: &str = "ai-setup-cli"; +const AI_SETUP_DEFAULT_VERSION: &str = "latest"; +const AI_SETUP_PACKAGE_ENV: &str = "PUP_SETUP_AI_SETUP_PACKAGE"; +const MIN_NODE_MAJOR: u32 = 22; +const NODE_DOWNLOAD_URL: &str = "https://nodejs.org/en/download"; +pub const LOGIN_EXTRA_SCOPES: &str = "api_keys_read"; +const TELEMETRY_TIMEOUT: Duration = Duration::from_secs(2); +const TELEMETRY_EVENT: &str = "setup_prerequisite_missing"; +const TELEMETRY_SERVICE: &str = "setup-cli"; +const TELEMETRY_TAGS: &str = "team:agentic-onboarding,service:setup-cli,source:pup"; + +#[derive(Debug, PartialEq, Eq)] +enum NodeStatus { + Ready(PathBuf), + NpxNotFound, + NodeNotFound, + NodeTooOld(String), +} + +impl NodeStatus { + fn reason(&self) -> &'static str { + match self { + NodeStatus::Ready(_) => "ready", + NodeStatus::NpxNotFound => "npx_not_found", + NodeStatus::NodeNotFound => "node_not_found", + NodeStatus::NodeTooOld(_) => "node_too_old", + } + } + + fn node_version(&self) -> Option<&str> { + match self { + NodeStatus::NodeTooOld(version) => Some(version), + _ => None, + } + } + + fn message(&self) -> String { + match self { + NodeStatus::NodeTooOld(version) => format!( + "pup setup needs Node.js {MIN_NODE_MAJOR} or newer (found {version}). \ + Upgrade from {NODE_DOWNLOAD_URL} and rerun `pup setup`." + ), + NodeStatus::NpxNotFound => format!( + "pup setup runs Datadog's AI Setup with npx, which wasn't found. npx comes \ + with npm: install npm (for example `sudo apt install npm`) or reinstall \ + Node.js from {NODE_DOWNLOAD_URL}, then rerun `pup setup`." + ), + _ => format!( + "pup setup runs Datadog's AI Setup with npx, which needs Node.js \ + {MIN_NODE_MAJOR} or newer. Install it from {NODE_DOWNLOAD_URL} and rerun \ + `pup setup`." + ), + } + } +} + +pub async fn run( + cfg: &mut Config, + args: &[String], + agent: bool, + login: impl AsyncFnOnce(&Config) -> Result, +) -> Result { + let args = &normalize_help_args(args); + let status = probe_node(); + let NodeStatus::Ready(npx) = status else { + return Ok(report_missing_node(cfg, &status).await); + }; + let package = package_spec(std::env::var(AI_SETUP_PACKAGE_ENV).ok().as_deref())?; + if asks_for_help(args) && !has_flag(args, "--json") { + if let Some(help) = fetch_ai_setup_help(&npx, &package, cfg) { + setup_help::print(&help)?; + return Ok(0); + } + } + let headless = asks_for_headless(args) + || wants_headless( + agent, + std::io::stdin().is_terminal(), + std::io::stdout().is_terminal(), + ); + let had_session = cfg.access_token.is_some(); + let mut login = Some(login); + if should_log_in(headless, had_session, args) { + let granted_scopes = stored_token_scopes(cfg); + if let Some(login) = login.take() { + ensure_session(cfg, granted_scopes.as_deref(), login, |cfg| { + crate::settings::load_token_from_storage(&cfg.site, cfg.org.as_deref()) + }) + .await?; + } + } + let launch = Launch { package, headless }; + let cmd = build_command(&npx, &build_args(args, &cfg.site, &launch), cfg); + let exit_code = exec::run_inherited(cmd, "npx")?; + if should_offer_pup_login(headless, had_session, args, exit_code) { + if let Some(login) = login.take() { + offer_pup_login(cfg, login).await; + } + } + Ok(exit_code) +} + +pub fn apply_site_arg(cfg: &mut Config, args: &[String]) -> Result<()> { + let site_flags = args + .iter() + .filter(|arg| *arg == "--site" || arg.starts_with("--site=")) + .count(); + if site_flags > 1 { + anyhow::bail!("--site was given more than once"); + } + let Some(site) = flag_value(args, "--site") else { + return Ok(()); + }; + let previous = cfg.site.clone(); + cfg.set_site_explicit(site)?; + let token_from_env = std::env::var("DD_ACCESS_TOKEN").is_ok_and(|token| !token.is_empty()); + if cfg.site != previous && !token_from_env { + cfg.access_token = crate::settings::load_token_from_storage(&cfg.site, cfg.org.as_deref()); + } + Ok(()) +} + +fn flag_value(args: &[String], flag: &str) -> Option { + let prefix = format!("{flag}="); + let mut iter = args.iter(); + while let Some(arg) = iter.next() { + if arg == flag { + return Some(iter.next().cloned().unwrap_or_default()); + } + if let Some(value) = arg.strip_prefix(&prefix) { + return Some(value.to_string()); + } + } + None +} + +fn without_flag(args: &[String], flag: &str) -> Vec { + let prefix = format!("{flag}="); + let mut kept = Vec::with_capacity(args.len()); + let mut iter = args.iter(); + while let Some(arg) = iter.next() { + if arg == flag { + iter.next(); + } else if !arg.starts_with(&prefix) { + kept.push(arg.clone()); + } + } + kept +} + +fn should_offer_pup_login( + headless: bool, + had_session: bool, + args: &[String], + exit_code: i32, +) -> bool { + !headless && !had_session && exit_code == 0 && !asks_for_help(args) +} + +async fn offer_pup_login(cfg: &Config, login: impl AsyncFnOnce(&Config) -> Result) { + eprint!( + "\nLog in to pup on {} too, so you can use its other commands? [Y/n] ", + cfg.site + ); + std::io::Write::flush(&mut std::io::stderr()).ok(); + let mut answer = String::new(); + let read = std::io::stdin().read_line(&mut answer); + if !accepts_login(read, &answer) { + eprintln!("You can run `pup auth login` any time."); + return; + } + if let Err(err) = login(cfg).await { + eprintln!("pup login failed ({err}). You can run `pup auth login` any time."); + } +} + +fn accepts_login(read: std::io::Result, answer: &str) -> bool { + matches!(read, Ok(bytes) if bytes > 0) + && matches!(answer.trim().to_lowercase().as_str(), "" | "y" | "yes") +} + +fn signup_hint(site: &str) -> String { + format!( + "pup setup needs a Datadog login. No account yet? Sign up at https://{}/signup, \ + then rerun `pup setup`.", + crate::settings::auth_host_for(site) + ) +} + +#[derive(Debug)] +struct Launch { + package: String, + headless: bool, +} + +async fn report_missing_node(cfg: &Config, status: &NodeStatus) -> i32 { + eprintln!("{}", status.message()); + report_prerequisite_missing(cfg, status).await; + 1 +} + +async fn ensure_session( + cfg: &mut Config, + granted_scopes: Option<&str>, + login: impl AsyncFnOnce(&Config) -> Result, + stored_token: impl FnOnce(&Config) -> Option, +) -> Result<()> { + if !needs_login(cfg, granted_scopes) { + return Ok(()); + } + if cfg.access_token.is_none() { + eprintln!("{}", signup_hint(&cfg.site)); + } + let session = login(cfg).await?; + cfg.site = session.site; + cfg.org = session.org; + cfg.access_token = stored_token(cfg); + if cfg.access_token.is_none() { + anyhow::bail!("login finished but no pup session was found; run `pup auth login`"); + } + Ok(()) +} + +fn has_credentials(cfg: &Config) -> bool { + cfg.access_token.is_some() || (cfg.api_key.is_some() && cfg.app_key.is_some()) +} + +fn needs_login(cfg: &Config, granted_scopes: Option<&str>) -> bool { + if cfg.access_token.is_none() { + return !has_credentials(cfg); + } + granted_scopes.is_some_and(|scopes| !has_scopes(scopes, LOGIN_EXTRA_SCOPES)) +} + +fn has_scopes(granted: &str, required: &str) -> bool { + let granted: Vec<&str> = granted.split_whitespace().collect(); + crate::settings::parse_scopes(required) + .iter() + .all(|scope| granted.contains(&scope.as_str())) +} + +fn stored_token_scopes(cfg: &Config) -> Option { + if std::env::var("DD_ACCESS_TOKEN").is_ok_and(|token| !token.is_empty()) { + return None; + } + let guard = crate::auth::storage::get_storage().ok()?; + let lock = guard.lock().ok()?; + let tokens = lock + .as_ref()? + .load_tokens(&cfg.site, cfg.org.as_deref()) + .ok()??; + active_token_scopes(cfg, tokens) +} + +fn active_token_scopes(cfg: &Config, saved: crate::auth::types::TokenSet) -> Option { + (cfg.access_token.as_deref() == Some(saved.access_token.as_str())).then_some(saved.scope) +} + +fn should_log_in(headless: bool, has_session: bool, args: &[String]) -> bool { + (headless || has_session) && !asks_for_help(args) +} + +fn normalize_help_args(args: &[String]) -> Vec { + let mut args = args.to_vec(); + if args.first().is_some_and(|arg| arg == "help") { + args[0] = "--help".to_string(); + } + args +} + +fn fetch_ai_setup_help(npx: &Path, package: &str, cfg: &Config) -> Option { + let args = ["-y", "--package", package, AI_SETUP_BIN, "--help", "--json"].map(String::from); + let output = build_command(npx, &args, cfg) + .stdin(std::process::Stdio::null()) + .output() + .ok()?; + if !output.status.success() { + return None; + } + setup_help::parse(&output.stdout) +} + +fn asks_for_help(args: &[String]) -> bool { + args.iter().any(|arg| arg == "--help" || arg == "-h") +} + +pub fn is_help_request(args: &[String]) -> bool { + asks_for_help(&normalize_help_args(args)) +} + +async fn report_prerequisite_missing(cfg: &Config, status: &NodeStatus) { + let Some(token) = intake_client_token(&cfg.site) else { + return; + }; + let url = format!("https://http-intake.logs.{}/v1/input/{token}", cfg.site); + let dd_oid = cfg + .access_token + .as_ref() + .and_then(|_| crate::auth::storage::find_session(cfg.org.as_deref())) + .filter(|session| session.site == cfg.site) + .and_then(|session| session.org_uuid); + send_intake_event(&url, &prerequisite_event(cfg, status, dd_oid.as_deref())).await; +} + +fn intake_client_token(site: &str) -> Option<&'static str> { + match site { + "datad0g.com" => Some("pub9a1f0104f5632c5840cca5a99a43c32a"), + "datadoghq.com" => Some("pub5b34c05a3abc178e4d2abddf11bb4f31"), + "us3.datadoghq.com" => Some("pub82dbd99bf6299230cee3685535d4660d"), + "us5.datadoghq.com" => Some("pubf9a4a5b93cc00e7f4d5628f05ba1e7f0"), + "datadoghq.eu" => Some("pub471980dfdfb12f9cf2ffdf83551c9cb7"), + "ap1.datadoghq.com" => Some("pubbbc10e19ab9fa2f7eb54250610e9bbad"), + "ap2.datadoghq.com" => Some("pube391858a603cb55a0e72a8cb9f9fb627"), + "uk1.datadoghq.com" => Some("pub90fe05f76003a728ed554865d39a88d1"), + _ => None, + } +} + +fn prerequisite_event( + cfg: &Config, + status: &NodeStatus, + dd_oid: Option<&str>, +) -> serde_json::Value { + serde_json::json!({ + "message": TELEMETRY_EVENT, + "dd_oid": dd_oid, + "service": TELEMETRY_SERVICE, + "source": "pup", + "reason": status.reason(), + "node_version": status.node_version(), + "pup_version": crate::version::VERSION, + "os": std::env::consts::OS, + "arch": std::env::consts::ARCH, + "site": cfg.site, + "authenticated": has_credentials(cfg), + }) +} + +async fn send_intake_event(url: &str, event: &serde_json::Value) { + let Ok(client) = reqwest::Client::builder() + .timeout(TELEMETRY_TIMEOUT) + .build() + else { + return; + }; + let _ = client + .post(url) + .query(&[("ddsource", "pup"), ("ddtags", TELEMETRY_TAGS)]) + .json(event) + .send() + .await; +} + +fn wants_headless(agent: bool, stdin_tty: bool, stdout_tty: bool) -> bool { + agent || !stdin_tty || !stdout_tty +} + +fn package_spec(override_value: Option<&str>) -> Result { + let Some(value) = override_value.map(str::trim).filter(|v| !v.is_empty()) else { + return Ok(format!("{AI_SETUP_PACKAGE}@{AI_SETUP_DEFAULT_VERSION}")); + }; + if value.starts_with('-') || value.chars().any(|c| c.is_whitespace() || c.is_control()) { + anyhow::bail!( + "{AI_SETUP_PACKAGE_ENV} must be a version, a dist-tag, or an npm package spec \ + such as a tarball path, without spaces or a leading '-'" + ); + } + let is_full_spec = value.contains('/') || value.contains(':') || value.ends_with(".tgz"); + if is_full_spec { + Ok(value.to_string()) + } else { + Ok(format!("{AI_SETUP_PACKAGE}@{value}")) + } +} + +fn build_args(user_args: &[String], site: &str, launch: &Launch) -> Vec { + let mut args = vec![ + "-y".to_string(), + "--package".to_string(), + launch.package.clone(), + AI_SETUP_BIN.to_string(), + ]; + if launch.headless { + args.push("--headless".to_string()); + } + args.push("--site".to_string()); + args.push(site.to_string()); + args.extend( + without_flag(user_args, "--site") + .into_iter() + .filter(|arg| !is_headless_flag(arg)), + ); + args +} + +fn asks_for_headless(args: &[String]) -> bool { + args.iter() + .any(|arg| arg == "--headless" || arg == "--headless=true") +} + +fn is_headless_flag(arg: &str) -> bool { + arg == "--headless" || arg == "--no-headless" || arg.starts_with("--headless=") +} + +fn has_flag(args: &[String], flag: &str) -> bool { + args.iter() + .any(|arg| arg == flag || arg.starts_with(&format!("{flag}="))) +} + +fn build_command(npx: &Path, args: &[String], cfg: &Config) -> Command { + let mut cmd = Command::new(npx); + cmd.args(args); + exec::inject_auth_env(&mut cmd, cfg); + if cfg.access_token.is_some() || !has_credentials(cfg) { + cmd.env_remove("DD_API_KEY"); + cmd.env_remove("DD_APP_KEY"); + } + cmd +} + +fn probe_node() -> NodeStatus { + let path = std::env::var_os("PATH").unwrap_or_default(); + let npx = find_on_path("npx", &path); + let node = find_on_path("node", &path); + let version = node.and_then(|node| { + let output = Command::new(node).arg("--version").output().ok()?; + output + .status + .success() + .then(|| String::from_utf8_lossy(&output.stdout).trim().to_string()) + }); + node_status(npx, version.as_deref()) +} + +fn node_status(npx: Option, node_version: Option<&str>) -> NodeStatus { + let Some(version) = node_version else { + return NodeStatus::NodeNotFound; + }; + match parse_node_major(version) { + Some(major) if major >= MIN_NODE_MAJOR => {} + _ => return NodeStatus::NodeTooOld(version.to_string()), + } + match npx { + Some(npx) => NodeStatus::Ready(npx), + None => NodeStatus::NpxNotFound, + } +} + +fn parse_node_major(version: &str) -> Option { + version + .trim() + .trim_start_matches('v') + .split('.') + .next()? + .parse() + .ok() +} + +fn find_on_path(name: &str, path: &std::ffi::OsStr) -> Option { + let candidates: &[&str] = if cfg!(windows) { + &["exe", "cmd"] + } else { + &[""] + }; + std::env::split_paths(path).find_map(|dir| { + candidates.iter().find_map(|ext| { + let candidate = if ext.is_empty() { + dir.join(name) + } else { + dir.join(format!("{name}.{ext}")) + }; + candidate.is_file().then_some(candidate) + }) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::*; + + fn args(values: &[&str]) -> Vec { + values.iter().map(|v| v.to_string()).collect() + } + + fn env_value(cmd: &Command, name: &str) -> Option> { + cmd.get_envs() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.map(|v| v.to_string_lossy().to_string())) + } + + #[test] + fn headless_for_agents_and_non_terminals_only() { + assert!(wants_headless(true, true, true)); + assert!(wants_headless(false, false, true)); + assert!(wants_headless(false, true, false)); + assert!(!wants_headless(false, true, true)); + } + + fn launch(headless: bool) -> Launch { + Launch { + package: package_spec(None).unwrap(), + headless, + } + } + + #[test] + fn build_args_injects_headless_and_site_before_user_args() { + assert_eq!( + build_args( + &args(&["--product", "linux"]), + "datadoghq.eu", + &launch(true) + ), + args(&[ + "-y", + "--package", + "@datadog/ai-setup-cli@latest", + "ai-setup-cli", + "--headless", + "--site", + "datadoghq.eu", + "--product", + "linux", + ]) + ); + } + + #[test] + fn build_args_interactive_skips_headless() { + let built = build_args( + &args(&["--product", "rum"]), + "datadoghq.com", + &launch(false), + ); + assert!(!built.contains(&"--headless".to_string())); + assert!(built.contains(&"--site".to_string())); + } + + #[test] + fn build_args_passes_pup_site_once_in_place_of_user_site() { + for user in [ + args(&[ + "--headless", + "--site", + "https://app.us5.datadoghq.com", + "--product", + "linux", + ]), + args(&[ + "--headless", + "--site=https://app.us5.datadoghq.com", + "--product", + "linux", + ]), + ] { + let built = build_args(&user, "us5.datadoghq.com", &launch(true)); + assert_eq!(built.iter().filter(|a| *a == "--headless").count(), 1); + assert_eq!(built.iter().filter(|a| a.starts_with("--site")).count(), 1); + assert!(built.contains(&"us5.datadoghq.com".to_string())); + assert!(!built.iter().any(|a| a.contains("https:"))); + assert!(built.ends_with(&args(&["--product", "linux"]))); + } + } + + #[tokio::test] + async fn site_arg_becomes_pup_site_and_keeps_env_token() { + let _lock = lock_env().await; + std::env::set_var("DD_ACCESS_TOKEN", "env-token"); + let mut cfg = test_config("http://unused"); + cfg.site = "datadoghq.com".into(); + cfg.access_token = Some("env-token".into()); + + apply_site_arg( + &mut cfg, + &args(&["--product", "apm", "--site=https://app.datadoghq.eu/"]), + ) + .unwrap(); + assert_eq!(cfg.site, "datadoghq.eu"); + assert!(cfg.site_explicit); + assert_eq!(cfg.access_token.as_deref(), Some("env-token")); + + for missing in [ + args(&["--site", " "]), + args(&["--site="]), + args(&["--product", "linux", "--site"]), + ] { + assert!(apply_site_arg(&mut cfg, &missing).is_err(), "{missing:?}"); + } + cleanup_env(); + } + + #[test] + fn site_arg_given_twice_is_rejected() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let mut cfg = test_config("http://unused"); + cfg.site = "datadoghq.com".into(); + for repeated in [ + args(&["--site", "datadoghq.com", "--site", "datadoghq.eu"]), + args(&["--site=datadoghq.com", "--site", "datadoghq.eu"]), + args(&["--site", "datadoghq.com", "--site"]), + ] { + assert!(apply_site_arg(&mut cfg, &repeated).is_err(), "{repeated:?}"); + assert_eq!(cfg.site, "datadoghq.com"); + } + cleanup_env(); + } + + #[test] + fn user_cannot_turn_off_required_headless() { + for user in [ + args(&["--headless=false", "--product", "linux"]), + args(&["--no-headless", "--product", "linux"]), + args(&["--headless", "--no-headless", "--product", "linux"]), + ] { + let built = build_args(&user, "datadoghq.com", &launch(true)); + assert_eq!(built.iter().filter(|a| is_headless_flag(a)).count(), 1); + assert!(built.contains(&"--headless".to_string())); + assert!(built.ends_with(&args(&["--product", "linux"]))); + } + } + + #[test] + fn only_true_headless_flags_ask_for_headless() { + assert!(asks_for_headless(&args(&["--headless"]))); + assert!(asks_for_headless(&args(&["--headless=true"]))); + assert!(!asks_for_headless(&args(&["--headless=false"]))); + assert!(!asks_for_headless(&args(&["--no-headless"]))); + let interactive = build_args( + &args(&["--headless=false", "--product", "rum"]), + "datadoghq.com", + &launch(false), + ); + assert!(!interactive.iter().any(|a| is_headless_flag(a))); + } + + #[test] + fn flag_value_reads_both_forms() { + assert_eq!( + flag_value( + &args(&["--product", "apm", "--site", "datadoghq.eu"]), + "--site" + ), + Some("datadoghq.eu".into()) + ); + assert_eq!( + flag_value(&args(&["--site=datadoghq.eu"]), "--site"), + Some("datadoghq.eu".into()) + ); + assert_eq!(flag_value(&args(&["--product", "apm"]), "--site"), None); + } + + #[test] + fn build_args_passes_help_through() { + let built = build_args(&args(&["--help"]), "datadoghq.com", &launch(false)); + assert_eq!(built.last().map(String::as_str), Some("--help")); + assert!(!built.contains(&"--json".to_string())); + } + + #[test] + fn package_spec_defaults_to_latest() { + assert_eq!(package_spec(None).unwrap(), "@datadog/ai-setup-cli@latest"); + assert_eq!( + package_spec(Some(" ")).unwrap(), + "@datadog/ai-setup-cli@latest" + ); + } + + #[test] + fn package_spec_accepts_versions_tags_and_local_builds() { + assert_eq!( + package_spec(Some("2.0.25")).unwrap(), + "@datadog/ai-setup-cli@2.0.25" + ); + assert_eq!( + package_spec(Some("next")).unwrap(), + "@datadog/ai-setup-cli@next" + ); + for spec in [ + "/tmp/datadog-ai-setup-cli-2.1.0.tgz", + "./datadog-ai-setup-cli-2.1.0.tgz", + "datadog-ai-setup-cli-2.1.0.tgz", + "file:../ai-setup-cli", + "@datadog/ai-setup-cli@2.1.0", + ] { + assert_eq!(package_spec(Some(spec)).unwrap(), spec); + } + } + + #[test] + fn package_spec_rejects_flag_injection_and_whitespace() { + for bad in ["--registry=https://evil.example", "-y", "2.0 --foo", "a\tb"] { + let err = package_spec(Some(bad)).unwrap_err().to_string(); + assert!(err.contains("PUP_SETUP_AI_SETUP_PACKAGE"), "got: {err}"); + } + } + + #[test] + fn build_command_strips_keys_when_oauth_token_present() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let mut cfg = test_config("http://unused"); + cfg.access_token = Some("oauth-token".into()); + let cmd = build_command(Path::new("/usr/bin/npx"), &[], &cfg); + assert_eq!(env_value(&cmd, "DD_API_KEY"), Some(None)); + assert_eq!(env_value(&cmd, "DD_APP_KEY"), Some(None)); + assert_eq!( + env_value(&cmd, "DD_ACCESS_TOKEN"), + Some(Some("oauth-token".into())) + ); + assert_eq!( + env_value(&cmd, "DD_SITE"), + Some(Some("datadoghq.com".into())) + ); + cleanup_env(); + } + + #[test] + fn build_command_keeps_keys_without_oauth_token() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let cfg = test_config("http://unused"); + let cmd = build_command(Path::new("/usr/bin/npx"), &[], &cfg); + assert_eq!( + env_value(&cmd, "DD_API_KEY"), + Some(Some("test-api-key".into())) + ); + assert_eq!(env_value(&cmd, "DD_ACCESS_TOKEN"), Some(None)); + cleanup_env(); + } + + #[test] + fn build_command_forwards_output_format() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let cfg = test_config("http://unused"); + let cmd = build_command(Path::new("/usr/bin/npx"), &args(&["--help"]), &cfg); + assert_eq!(env_value(&cmd, "PUP_OUTPUT"), Some(Some("json".into()))); + assert!(cmd.get_args().any(|a| a == "--help")); + cleanup_env(); + } + + #[test] + fn parse_node_major_handles_versions() { + assert_eq!(parse_node_major("v22.3.0\n"), Some(22)); + assert_eq!(parse_node_major("18.19.0"), Some(18)); + assert_eq!(parse_node_major("garbage"), None); + assert_eq!(parse_node_major(""), None); + } + + #[test] + fn missing_node_messages_match_the_cause() { + let npx = NodeStatus::NpxNotFound.message(); + assert!(npx.contains("npx, which wasn't found"), "{npx}"); + assert!(npx.contains("install npm"), "{npx}"); + let missing = NodeStatus::NodeNotFound.message(); + assert!(missing.contains("needs Node.js 22 or newer"), "{missing}"); + assert!(!missing.contains("install npm"), "{missing}"); + let old = NodeStatus::NodeTooOld("v20.20.1".into()).message(); + assert!(old.contains("found v20.20.1"), "{old}"); + } + + #[test] + fn node_status_cases() { + let npx = PathBuf::from("/bin/npx"); + assert_eq!( + node_status(Some(npx.clone()), Some("v22.1.0")), + NodeStatus::Ready(npx.clone()) + ); + assert_eq!(node_status(None, Some("v24.0.0")), NodeStatus::NpxNotFound); + assert_eq!( + node_status(Some(npx.clone()), None), + NodeStatus::NodeNotFound + ); + assert_eq!( + node_status(Some(npx), Some("v18.19.0")), + NodeStatus::NodeTooOld("v18.19.0".into()) + ); + let too_old = NodeStatus::NodeTooOld("v18.19.0".into()); + assert_eq!(too_old.reason(), "node_too_old"); + assert_eq!(too_old.node_version(), Some("v18.19.0")); + assert_eq!(NodeStatus::NpxNotFound.node_version(), None); + } + + fn no_credentials(cfg: &mut Config) { + cfg.api_key = None; + cfg.app_key = None; + cfg.access_token = None; + } + + fn saved_session(cfg: &Config) -> SessionEntry { + SessionEntry { + site: cfg.site.clone(), + org: cfg.org.clone(), + org_uuid: None, + } + } + + #[tokio::test] + async fn no_credentials_logs_in_then_uses_new_token() { + let _lock = lock_env().await; + let mut cfg = test_config("http://unused"); + no_credentials(&mut cfg); + let called = std::cell::Cell::new(false); + let result = ensure_session( + &mut cfg, + None, + async |cfg: &Config| { + called.set(true); + Ok(saved_session(cfg)) + }, + |_| Some("new-token".into()), + ) + .await; + assert!(result.is_ok()); + assert!(called.get()); + assert_eq!(cfg.access_token.as_deref(), Some("new-token")); + cleanup_env(); + } + + #[tokio::test] + async fn session_login_decision_follows_granted_scopes() { + let _lock = lock_env().await; + let cases = [ + (None, false), + (Some("apm_read api_keys_read dashboards_read"), false), + (Some("apm_read dashboards_read"), true), + (Some(""), true), + ]; + for (granted_scopes, expect_login) in cases { + let mut cfg = test_config("http://unused"); + cfg.access_token = Some("old-token".into()); + let called = std::cell::Cell::new(false); + let result = ensure_session( + &mut cfg, + granted_scopes, + async |cfg: &Config| { + called.set(true); + Ok(saved_session(cfg)) + }, + |_| Some("new-token".into()), + ) + .await; + assert!(result.is_ok(), "{granted_scopes:?}"); + assert_eq!(called.get(), expect_login, "{granted_scopes:?}"); + let expected_token = if expect_login { + "new-token" + } else { + "old-token" + }; + assert_eq!(cfg.access_token.as_deref(), Some(expected_token)); + } + cleanup_env(); + } + + #[test] + fn key_pair_without_session_skips_login() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let mut cfg = test_config("http://unused"); + cfg.access_token = None; + cfg.api_key = Some("api".into()); + cfg.app_key = Some("app".into()); + assert!(!needs_login(&cfg, None)); + cleanup_env(); + } + + #[test] + fn saved_scopes_count_only_for_the_active_token() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let saved = |token: &str| crate::auth::types::TokenSet { + access_token: token.into(), + refresh_token: String::new(), + token_type: "Bearer".into(), + expires_in: 3600, + issued_at: 0, + scope: "apm_read".into(), + client_id: String::new(), + }; + let mut cfg = test_config("http://unused"); + cfg.access_token = Some("config-token".into()); + assert_eq!(active_token_scopes(&cfg, saved("stored-token")), None); + assert_eq!( + active_token_scopes(&cfg, saved("config-token")).as_deref(), + Some("apm_read") + ); + cleanup_env(); + } + + #[tokio::test] + async fn login_failure_stops_setup() { + let _lock = lock_env().await; + let mut cfg = test_config("http://unused"); + no_credentials(&mut cfg); + let result = ensure_session( + &mut cfg, + None, + async |_: &Config| Err(anyhow::anyhow!("login cancelled")), + |_| Some("unused".into()), + ) + .await; + assert_eq!(result.unwrap_err().to_string(), "login cancelled"); + assert!(cfg.access_token.is_none()); + cleanup_env(); + } + + #[tokio::test] + async fn login_without_saved_session_is_an_error() { + let _lock = lock_env().await; + let mut cfg = test_config("http://unused"); + no_credentials(&mut cfg); + let result = ensure_session( + &mut cfg, + None, + async |cfg: &Config| Ok(saved_session(cfg)), + |_| None, + ) + .await; + assert!(result.unwrap_err().to_string().contains("no pup session")); + cleanup_env(); + } + + #[tokio::test] + async fn login_that_switches_region_or_org_loads_that_session() { + let _lock = lock_env().await; + let mut cfg = test_config("http://unused"); + no_credentials(&mut cfg); + cfg.site = "datadoghq.com".into(); + let result = ensure_session( + &mut cfg, + None, + async |_: &Config| { + Ok(SessionEntry { + site: "us3.datadoghq.com".into(), + org: Some("acme".into()), + org_uuid: None, + }) + }, + |cfg| { + (cfg.site == "us3.datadoghq.com" && cfg.org.as_deref() == Some("acme")) + .then(|| "us3-token".to_string()) + }, + ) + .await; + assert!(result.is_ok()); + assert_eq!(cfg.site, "us3.datadoghq.com"); + assert_eq!(cfg.access_token.as_deref(), Some("us3-token")); + cleanup_env(); + } + + #[test] + fn leading_help_word_is_treated_as_help_flag() { + assert_eq!( + normalize_help_args(&args(&["help", "--product", "linux"])), + args(&["--help", "--product", "linux"]) + ); + assert!(asks_for_help(&normalize_help_args(&args(&["help"])))); + assert_eq!( + normalize_help_args(&args(&["--product", "help"])), + args(&["--product", "help"]) + ); + } + + #[test] + fn offers_pup_login_only_after_successful_terminal_runs_without_a_session() { + let setup = args(&["--product", "apm"]); + assert!(should_offer_pup_login(false, false, &setup, 0)); + assert!(!should_offer_pup_login(true, false, &setup, 0)); + assert!(!should_offer_pup_login(false, true, &setup, 0)); + assert!(!should_offer_pup_login(false, false, &setup, 1)); + assert!(!should_offer_pup_login(false, false, &args(&["--help"]), 0)); + } + + #[test] + fn pup_login_offer_defaults_to_yes() { + for answer in ["\n", "y\n", "YES\n"] { + assert!(accepts_login(Ok(answer.len()), answer), "{answer:?}"); + } + for answer in ["n\n", "no\n", "later\n"] { + assert!(!accepts_login(Ok(answer.len()), answer), "{answer:?}"); + } + } + + #[test] + fn pup_login_offer_declines_on_end_of_input_or_read_error() { + assert!(!accepts_login(Ok(0), "")); + assert!(!accepts_login(Err(std::io::Error::other("closed")), "")); + } + + #[test] + fn signup_hint_points_at_the_site_ui_host() { + assert!(signup_hint("datadoghq.com").contains("https://app.datadoghq.com/signup")); + assert!(signup_hint("us5.datadoghq.com").contains("https://us5.datadoghq.com/signup")); + } + + #[test] + fn logs_in_for_headless_runs_or_existing_sessions_but_not_help() { + let setup = args(&["--product", "apm"]); + assert!(should_log_in(true, false, &setup)); + assert!(should_log_in(false, true, &setup)); + assert!(!should_log_in(false, false, &setup)); + assert!(!should_log_in(true, true, &args(&["--help"]))); + assert!(!should_log_in( + true, + false, + &args(&["--product", "apm", "-h"]) + )); + } + + #[test] + fn interactive_without_session_forwards_no_credentials() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let mut cfg = test_config("http://unused"); + no_credentials(&mut cfg); + let built = build_args( + &args(&["--product", "rum"]), + &cfg.site, + &launch(wants_headless(false, true, true)), + ); + assert!(!built.contains(&"--headless".to_string())); + let cmd = build_command(Path::new("/usr/bin/npx"), &built, &cfg); + for name in ["DD_ACCESS_TOKEN", "DD_API_KEY", "DD_APP_KEY"] { + assert_eq!(env_value(&cmd, name), Some(None), "{name} must be removed"); + } + cleanup_env(); + } + + #[test] + fn lone_api_key_is_not_forwarded() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let mut cfg = test_config("http://unused"); + cfg.app_key = None; + let cmd = build_command(Path::new("/usr/bin/npx"), &[], &cfg); + assert_eq!(env_value(&cmd, "DD_API_KEY"), Some(None)); + assert_eq!(env_value(&cmd, "DD_APP_KEY"), Some(None)); + cleanup_env(); + } + + #[test] + fn find_on_path_finds_files_only() { + let dir = std::env::temp_dir().join(format!("pup-setup-path-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(dir.join("npx-dir")).unwrap(); + let name = if cfg!(windows) { "node.exe" } else { "node" }; + std::fs::write(dir.join(name), "").unwrap(); + let path = std::env::join_paths([dir.clone()]).unwrap(); + assert_eq!(find_on_path("node", &path), Some(dir.join(name))); + assert_eq!(find_on_path("npx-dir", &path), None); + assert_eq!(find_on_path("npx", &path), None); + std::fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn intake_tokens_cover_datadog_sites_only() { + for site in [ + "datad0g.com", + "datadoghq.com", + "us3.datadoghq.com", + "us5.datadoghq.com", + "datadoghq.eu", + "ap1.datadoghq.com", + "ap2.datadoghq.com", + "uk1.datadoghq.com", + ] { + assert!( + intake_client_token(site).is_some_and(|t| t.starts_with("pub")), + "{site}" + ); + } + assert_eq!(intake_client_token("mygateway.example.com"), None); + } + + #[test] + fn prerequisite_event_describes_the_failure() { + let _guard = crate::test_utils::ENV_LOCK.blocking_lock(); + let mut cfg = test_config("http://unused"); + cfg.site = "datad0g.com".into(); + no_credentials(&mut cfg); + let event = prerequisite_event(&cfg, &NodeStatus::NodeTooOld("v18.19.0".into()), None); + assert_eq!(event["message"], "setup_prerequisite_missing"); + assert_eq!(event["service"], "setup-cli"); + assert_eq!(event["source"], "pup"); + assert_eq!(event["reason"], "node_too_old"); + assert_eq!(event["node_version"], "v18.19.0"); + assert_eq!(event["site"], "datad0g.com"); + assert_eq!(event["authenticated"], false); + assert_eq!(event["pup_version"], crate::version::VERSION); + assert!(event["dd_oid"].is_null()); + + let org_uuid = "1dd437f2-ecb6-11f0-8e9a-0217575d1df7"; + let event = prerequisite_event(&cfg, &NodeStatus::NodeNotFound, Some(org_uuid)); + assert_eq!(event["dd_oid"], org_uuid); + cleanup_env(); + } + + #[tokio::test] + async fn intake_event_is_posted_with_setup_cli_tags() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("POST", "/v1/input/pub-test") + .match_query(mockito::Matcher::AllOf(vec![ + mockito::Matcher::UrlEncoded("ddsource".into(), "pup".into()), + mockito::Matcher::UrlEncoded("ddtags".into(), TELEMETRY_TAGS.into()), + ])) + .match_body(mockito::Matcher::PartialJson(serde_json::json!({ + "message": "setup_prerequisite_missing", + "reason": "node_not_found", + }))) + .with_status(202) + .create_async() + .await; + let cfg = test_config("http://unused"); + let event = prerequisite_event(&cfg, &NodeStatus::NodeNotFound, None); + send_intake_event(&format!("{}/v1/input/pub-test", server.url()), &event).await; + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn intake_failure_is_ignored() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("POST", mockito::Matcher::Any) + .with_status(500) + .create_async() + .await; + let cfg = test_config("http://unused"); + let event = prerequisite_event(&cfg, &NodeStatus::NpxNotFound, None); + send_intake_event(&format!("{}/v1/input/pub-test", server.url()), &event).await; + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn missing_node_exits_1_on_hosts_without_an_intake_token() { + let mut cfg = test_config("http://unused"); + cfg.site = "mygateway.example.com".into(); + assert_eq!( + report_missing_node(&cfg, &NodeStatus::NodeNotFound).await, + 1 + ); + } +} diff --git a/src/commands/setup_help.rs b/src/commands/setup_help.rs new file mode 100644 index 00000000..c5e08118 --- /dev/null +++ b/src/commands/setup_help.rs @@ -0,0 +1,172 @@ +use anyhow::Result; +use clap::{Arg, ArgAction}; +use serde::Deserialize; + +const PUP_MANAGED_FLAGS: &[&str] = &["--site", "--headless"]; +const USAGE: &str = "pup setup --product [OPTIONS] [PRODUCT FLAGS]"; + +#[derive(Debug, Deserialize)] +pub struct AiSetupHelp { + global_flags: Vec, + products: Vec, +} + +#[derive(Debug, Deserialize)] +struct HelpFlag { + name: String, + #[serde(rename = "type")] + kind: String, + #[serde(default)] + description: String, +} + +#[derive(Debug, Deserialize)] +struct HelpProduct { + name: String, + #[serde(default)] + display_name: String, + #[serde(default)] + description: String, + #[serde(default)] + flags: Vec, +} + +pub fn parse(json: &[u8]) -> Option { + serde_json::from_slice(json).ok() +} + +pub fn print(help: &AiSetupHelp) -> Result<()> { + let mut root = crate::cli_command(); + root.build(); + let base = root + .find_subcommand("setup") + .ok_or_else(|| anyhow::anyhow!("setup command is not registered"))?; + help_command(base, help).print_long_help()?; + Ok(()) +} + +fn help_command(base: &clap::Command, help: &AiSetupHelp) -> clap::Command { + let mut cmd = clap::Command::new("setup") + .override_usage(USAGE) + .after_long_help(products_section(&help.products)); + if let Some(about) = base.get_about() { + cmd = cmd.about(about.clone()); + } + if let Some(long_about) = base.get_long_about() { + cmd = cmd.long_about(long_about.clone()); + } + cmd.args( + help.global_flags + .iter() + .filter(|flag| !PUP_MANAGED_FLAGS.contains(&flag.name.as_str())) + .filter_map(flag_arg), + ) + .args( + base.get_arguments() + .filter(|arg| arg.is_global_set()) + .cloned(), + ) +} + +fn products_section(products: &[HelpProduct]) -> String { + let name_width = products.iter().map(|p| p.name.len()).max().unwrap_or(0); + let mut out = String::from("Products:\n"); + for product in products { + let summary = match ( + product.display_name.is_empty(), + product.description.is_empty(), + ) { + (false, false) => format!("{}: {}", product.display_name, product.description), + (false, true) => product.display_name.clone(), + _ => product.description.clone(), + }; + out.push_str(&format!(" {: = product.flags.iter().map(flag_label).collect(); + let flag_width = labels.iter().map(String::len).max().unwrap_or(0); + for (label, flag) in labels.iter().zip(&product.flags) { + out.push_str(&format!( + " {label: Option { + let long = flag.name.strip_prefix("--")?.to_string(); + let arg = Arg::new(long.clone()) + .long(long) + .help(flag.description.clone()); + Some(match value_name(flag) { + Some(value) => arg.value_name(value).action(ArgAction::Set), + None => arg.action(ArgAction::SetTrue), + }) +} + +fn flag_label(flag: &HelpFlag) -> String { + match value_name(flag) { + Some(value) => format!("{} <{value}>", flag.name), + None => flag.name.clone(), + } +} + +fn value_name(flag: &HelpFlag) -> Option { + (flag.kind != "boolean").then(|| { + flag.name + .trim_start_matches('-') + .to_uppercase() + .replace('-', "_") + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + const SAMPLE: &str = r#"{ + "name": "ai-setup-cli", + "version": "1.0.0", + "usage": "ai-setup-cli --site --product [product flags]", + "global_flags": [ + {"name": "--site", "type": "string", "required": true, "description": "Datadog site"}, + {"name": "--product", "type": "string", "required": false, "required_in_headless": true, "description": "Product to set up (see products)"}, + {"name": "--headless", "type": "boolean", "required": false, "description": "Run setup non-interactively"}, + {"name": "--auto-approve", "type": "boolean", "required": false, "description": "Auto-approve all bash commands"} + ], + "products": [ + {"name": "linux", "display_name": "Linux Observability", "description": "Metrics, logs and traces from Linux", + "flags": [{"name": "--env", "type": "string", "required": false, "description": "Environment name"}]}, + {"name": "otel", "display_name": "Datadog OpenTelemetry", "description": "Direct OTLP export", "flags": []} + ] + }"#; + + fn sample() -> AiSetupHelp { + parse(SAMPLE.as_bytes()).expect("sample help parses") + } + + fn built_setup() -> clap::Command { + let mut root = crate::cli_command(); + root.build(); + help_command(root.find_subcommand("setup").unwrap(), &sample()) + } + + #[test] + fn parse_rejects_text_help() { + assert!(parse(b"Interactive tool for Datadog setup").is_none()); + } + + #[test] + fn text_help_uses_pup_usage_and_lists_products() { + let mut cmd = built_setup(); + let text = cmd.render_long_help().to_string(); + assert!(text.contains(&format!("Usage: {USAGE}")), "{text}"); + assert!(text.contains("--product "), "{text}"); + assert!(text.contains("--read-only"), "{text}"); + assert!(text.contains("Products:"), "{text}"); + assert!(text.contains("linux Linux Observability: Metrics, logs and traces from Linux")); + assert!(text.contains("--env Environment name"), "{text}"); + assert!(!text.contains("--headless"), "{text}"); + assert!(!text.contains("ai-setup-cli"), "{text}"); + } +} diff --git a/src/extensions/exec.rs b/src/extensions/exec.rs index 87d84161..8a57b1d8 100644 --- a/src/extensions/exec.rs +++ b/src/extensions/exec.rs @@ -11,12 +11,16 @@ pub fn exec_extension(ext_path: &Path, args: &[String], cfg: &Config) -> Result< inject_auth_env(&mut cmd, cfg); + run_inherited(cmd, &format!("extension {}", ext_path.display())) +} + +pub(crate) fn run_inherited(mut cmd: std::process::Command, label: &str) -> Result { let status = cmd .stdin(std::process::Stdio::inherit()) .stdout(std::process::Stdio::inherit()) .stderr(std::process::Stdio::inherit()) .status() - .map_err(|e| anyhow::anyhow!("failed to execute extension {}: {e}", ext_path.display()))?; + .map_err(|e| anyhow::anyhow!("failed to execute {label}: {e}"))?; // On Unix, if the process was killed by a signal, status.code() returns None. // Use the standard convention of 128 + signal_number. @@ -37,7 +41,7 @@ pub fn exec_extension(ext_path: &Path, args: &[String], cfg: &Config) -> Result< /// Set (or remove) auth and config environment variables on the child process command. /// Variables not active in the current config are explicitly removed to prevent /// stale credentials from leaking through the parent environment. -fn inject_auth_env(cmd: &mut std::process::Command, cfg: &Config) { +pub(crate) fn inject_auth_env(cmd: &mut std::process::Command, cfg: &Config) { // Always set site and output format. cmd.env("DD_SITE", &cfg.site); cmd.env("PUP_OUTPUT", cfg.output_format.to_string()); diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs index 04f1285c..2440787d 100644 --- a/src/extensions/mod.rs +++ b/src/extensions/mod.rs @@ -322,6 +322,12 @@ mod tests { assert!(!is_builtin_command("terraform")); } + #[test] + fn test_is_builtin_setup() { + assert!(is_builtin_command("setup")); + assert!(!is_builtin_command("setupx")); + } + #[test] fn test_is_builtin_extension() { assert!(is_builtin_command("extension")); diff --git a/src/main.rs b/src/main.rs index 6e88cb3f..e0e67854 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1461,6 +1461,21 @@ enum Commands { #[command(subcommand)] action: ExperimentActions, }, + /// Set up Datadog products in your project. Run `pup setup --help` to list products and their flags. + /// + /// Runs Datadog's AI Setup CLI from npm with npx, so it needs Node.js 22 + /// or newer. Every argument after `setup` is passed through to AI Setup. + /// + /// EXAMPLES: + /// pup setup --product apm + /// pup setup --help + #[cfg(not(target_arch = "wasm32"))] + #[command(verbatim_doc_comment, disable_help_flag = true)] + Setup { + /// Arguments passed through to AI Setup + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] + args: Vec, + }, /// Manage feature flags /// /// Manage Datadog feature flags and their environments. @@ -12621,6 +12636,104 @@ enum AuthActions { Test, } +fn setup_help_alias(mut args: Vec) -> Vec { + let globals = global_flag_specs(); + let Some(help_at) = first_positional_index(&args, &globals) else { + return args; + }; + if args[help_at] != "help" || args.get(help_at + 1).map(String::as_str) != Some("setup") { + return args; + } + args.remove(help_at); + args.push("--help".to_string()); + args +} + +fn hoist_setup_global_flags(args: Vec) -> Vec { + if !args.iter().any(|arg| arg == "setup") { + return args; + } + let globals = global_flag_specs(); + let Some(setup_at) = first_positional_index(&args, &globals).filter(|&at| args[at] == "setup") + else { + return args; + }; + let mut hoisted = args[..setup_at].to_vec(); + let mut passthrough = Vec::new(); + let mut rest = args[setup_at + 1..].iter(); + while let Some(token) = rest.next() { + if token == "--" { + passthrough.push(token.clone()); + passthrough.extend(rest.by_ref().cloned()); + break; + } + match match_global_flag(token, &globals) { + Some(takes_value) => { + hoisted.push(token.clone()); + if takes_value && !token.contains('=') { + hoisted.extend(rest.next().cloned()); + } + } + None => passthrough.push(token.clone()), + } + } + hoisted.push("setup".to_string()); + hoisted.extend(passthrough); + hoisted +} + +struct GlobalFlag { + long: Option, + short: Option, + takes_value: bool, +} + +fn global_flag_specs() -> Vec { + cli_command() + .get_arguments() + .filter(|arg| arg.is_global_set()) + .map(|arg| GlobalFlag { + long: arg.get_long().map(str::to_string), + short: arg.get_short(), + takes_value: arg.get_action().takes_values(), + }) + .collect() +} + +fn match_global_flag(token: &str, globals: &[GlobalFlag]) -> Option { + globals + .iter() + .find(|flag| { + let long_match = flag.long.as_deref().is_some_and(|long| { + token + .strip_prefix("--") + .is_some_and(|rest| rest == long || rest.starts_with(&format!("{long}="))) + }); + let short_match = flag.short.is_some_and(|short| token == format!("-{short}")); + long_match || short_match + }) + .map(|flag| flag.takes_value) +} + +fn first_positional_index(args: &[String], globals: &[GlobalFlag]) -> Option { + let mut index = 1; + while index < args.len() { + let token = &args[index]; + match match_global_flag(token, globals) { + Some(takes_value) => { + index += if takes_value && !token.contains('=') { + 2 + } else { + 1 + }; + } + None if token.starts_with('-') => index += 1, + None => return Some(index), + } + } + None +} + // ---- Agent JSON schema (`pup agent schema`) ---- /// Query syntax hints per domain, shared by the full and scoped agent schemas. @@ -13087,6 +13200,7 @@ pub(crate) fn is_write_command_name(name: &str) -> bool { || name == "pause" || name == "resume" || name == "generate" + || name == "setup" } fn build_command_schema(cmd: &clap::Command, parent_path: &str) -> serde_json::Value { @@ -13359,6 +13473,198 @@ mod test_agent_schema { } } + #[test] + fn setup_is_listed_in_agent_schema_as_a_write_command() { + let schema = get_schema(); + let commands = schema["commands"].as_array().unwrap(); + let setup = find_command(commands, &["setup"]).expect("setup must be in the schema"); + assert_eq!(setup["read_only"], serde_json::json!(false)); + let description = setup["description"].as_str().unwrap_or_default(); + assert!(description.contains("Set up Datadog products")); + } + + #[test] + fn pup_flags_after_setup_are_handled_by_pup() { + let argv = |items: &[&str]| items.iter().map(|s| s.to_string()).collect::>(); + let cases: &[(&[&str], &[&str])] = &[ + ( + &["pup", "setup", "--no-agent", "--product", "linux"], + &["pup", "--no-agent", "setup", "--product", "linux"], + ), + ( + &[ + "pup", + "setup", + "--product", + "linux", + "--org", + "eu", + "--read-only", + ], + &[ + "pup", + "--org", + "eu", + "--read-only", + "setup", + "--product", + "linux", + ], + ), + ( + &[ + "pup", "--agent", "setup", "--org=eu", "-o", "json", "--env", "prod", + ], + &[ + "pup", "--agent", "--org=eu", "-o", "json", "setup", "--env", "prod", + ], + ), + ( + &["pup", "setup", "--product", "linux", "--", "--agent"], + &["pup", "setup", "--product", "linux", "--", "--agent"], + ), + ( + &["pup", "monitors", "list", "--no-agent"], + &["pup", "monitors", "list", "--no-agent"], + ), + ( + &["pup", "--org", "setup", "monitors", "list"], + &["pup", "--org", "setup", "monitors", "list"], + ), + ]; + for (input, expected) in cases { + assert_eq!( + hoist_setup_global_flags(argv(input)), + argv(expected), + "{input:?}" + ); + } + } + + #[test] + fn pup_flags_in_and_after_a_setup_alias_are_handled_by_pup() { + let aliases = std::collections::BTreeMap::from([( + "install-linux".to_string(), + "setup --product linux --read-only".to_string(), + )]); + let args = hoist_setup_global_flags(commands::alias::apply_expansion( + &["pup", "install-linux", "--org", "prod"].map(String::from), + "install-linux", + &aliases, + )); + let cli = Cli::try_parse_from(args).unwrap(); + assert!(cli.read_only); + assert_eq!(cli.org.as_deref(), Some("prod")); + match cli.command { + Commands::Setup { args } => assert_eq!(args, ["--product", "linux"]), + _ => panic!("expected the setup command"), + } + } + + #[test] + fn read_only_after_setup_still_blocks_it() { + let args = hoist_setup_global_flags( + ["pup", "setup", "--product", "linux", "--read-only"] + .iter() + .map(|s| s.to_string()) + .collect(), + ); + let cli = Cli::try_parse_from(args).unwrap(); + assert!(cli.read_only); + match cli.command { + Commands::Setup { args } => assert_eq!(args, ["--product", "linux"]), + _ => panic!("expected the setup command"), + } + } + + #[test] + fn read_only_allows_setup_help_but_not_setup() { + let exempt = |input: &[&str]| { + let matches = cli_command().get_matches_from(input); + is_read_only_exempt(&matches) + }; + assert!(exempt(&["pup", "--read-only", "setup", "--help"])); + assert!(exempt(&["pup", "--read-only", "setup", "help"])); + assert!(exempt(&[ + "pup", + "--read-only", + "setup", + "--product", + "apm", + "-h" + ])); + assert!(!exempt(&[ + "pup", + "--read-only", + "setup", + "--product", + "apm" + ])); + assert!(!exempt(&["pup", "--read-only", "setup"])); + } + + #[test] + fn help_setup_is_an_alias_for_setup_help() { + let argv = |items: &[&str]| items.iter().map(|s| s.to_string()).collect::>(); + assert_eq!( + setup_help_alias(argv(&["pup", "help", "setup"])), + argv(&["pup", "setup", "--help"]) + ); + assert_eq!( + setup_help_alias(argv(&[ + "pup", + "--agent", + "help", + "setup", + "--product", + "linux" + ])), + argv(&["pup", "--agent", "setup", "--product", "linux", "--help"]) + ); + assert_eq!( + setup_help_alias(argv(&["pup", "help", "monitors"])), + argv(&["pup", "help", "monitors"]) + ); + } + + #[test] + fn setup_passes_all_arguments_through() { + let cli = + Cli::try_parse_from(["pup", "setup", "--product", "apm", "--help", "-x"]).unwrap(); + match cli.command { + Commands::Setup { args } => assert_eq!(args, ["--product", "apm", "--help", "-x"]), + _ => panic!("expected the setup command"), + } + } + + #[test] + fn setup_description_points_to_its_help() { + let schema = get_schema(); + let commands = schema["commands"].as_array().unwrap(); + let setup = find_command(commands, &["setup"]).unwrap(); + let description = setup["description"].as_str().unwrap_or_default(); + assert!( + description.contains("Run `pup setup --help` to list products and their flags"), + "got: {description}" + ); + } + + #[test] + fn setup_accepts_no_arguments() { + let cli = Cli::try_parse_from(["pup", "setup"]).unwrap(); + match cli.command { + Commands::Setup { args } => assert!(args.is_empty()), + _ => panic!("expected the setup command"), + } + } + + #[test] + fn setup_is_a_write_command_but_similar_names_are_not() { + assert!(is_write_command_name("setup")); + assert!(!is_write_command_name("setups")); + assert!(!is_write_command_name("list")); + } + #[test] fn schema_has_commands_array() { let schema = get_schema(); @@ -14383,6 +14689,13 @@ pub(crate) fn is_read_only_exempt(matches: &clap::ArgMatches) -> bool { matches.subcommand().and_then(|(_, m)| m.subcommand()), Some(("remote", _)) ), + #[cfg(not(target_arch = "wasm32"))] + Some("setup") => matches + .subcommand_matches("setup") + .and_then(|m| m.get_many::("args")) + .is_some_and(|args| { + commands::setup::is_help_request(&args.cloned().collect::>()) + }), _ => false, } } @@ -14943,7 +15256,7 @@ mod resolve_output_format_tests { } async fn main_inner() -> anyhow::Result<()> { - let args: Vec = std::env::args().collect(); + let args = hoist_setup_global_flags(setup_help_alias(std::env::args().collect())); // --- Extension interception (before clap parsing) --- // If the first positional arg is not a built-in command but matches an @@ -14980,7 +15293,7 @@ async fn main_inner() -> anyhow::Result<()> { let args = { let parsed = extensions::parse_extension_args(&args); if let Some(ref candidate) = parsed.candidate { - commands::alias::expand(&args, candidate) + hoist_setup_global_flags(commands::alias::expand(&args, candidate)) } else { args } @@ -15103,6 +15416,11 @@ async fn main_inner() -> anyhow::Result<()> { #[cfg(not(feature = "browser"))] let trusted_sites = settings::configured_trusted_sites(); + #[cfg(not(target_arch = "wasm32"))] + if let Commands::Setup { args } = &cli.command { + commands::setup::apply_site_arg(&mut cfg, args)?; + } + // Gate credential dispatch to non-Datadog hosts before any command dispatches. // Auth subcommands resolve their own site late (after --site is applied in-arm), // so they gate themselves; skip here to avoid a double prompt. @@ -19775,7 +20093,7 @@ async fn main_inner() -> anyhow::Result<()> { // Coerce empty `--org-uuid ""` to no-hint so callees treat // it the same as omitting the flag, not as `dd_oid=`. let org_uuid_hint = org_uuid.as_deref().filter(|s| !s.is_empty()); - commands::auth::login(&cfg, resolved, resolved_port, org_uuid_hint).await? + commands::auth::login(&cfg, resolved, resolved_port, org_uuid_hint).await?; } AuthActions::Logout => commands::auth::logout(&cfg).await?, AuthActions::Status { site } => { @@ -20694,6 +21012,23 @@ async fn main_inner() -> anyhow::Result<()> { } } } + #[cfg(not(target_arch = "wasm32"))] + Commands::Setup { args } => { + let login = + async |cfg: &config::Config| -> anyhow::Result { + let scopes = resolve_login_scopes( + None, + Some(commands::setup::LOGIN_EXTRA_SCOPES), + cfg.org.as_deref(), + cfg.read_only, + ); + let port = resolve_callback_port(None)?; + commands::auth::login(cfg, scopes, port, None).await + }; + let agent = useragent::agent_auto_approves(cli.no_agent); + let exit_code = commands::setup::run(&mut cfg, &args, agent, login).await?; + std::process::exit(exit_code); + } // --- Extensions --- #[cfg(not(target_arch = "wasm32"))] Commands::Extension { action } => match action {