From b5060c5eb4c511f320c9471e061e81bd58e83473 Mon Sep 17 00:00:00 2001 From: Teddy Kahwaji Date: Tue, 13 Oct 2026 12:35:30 -0400 Subject: [PATCH] feat(setup): Add AI-CLI Pup Setup command `pup setup` runs Datadog's AI Setup CLI with `npx -y --package @datadog/ai-setup-cli@latest ai-setup-cli`, passing every argument through. AI Setup is already published to npm, so this needs no extension install, release repository or checksums, at the cost of requiring Node.js 22+ when `setup` runs. - Checks for Node.js 22+ and npx, prints how to install what is missing, and sends one best-effort log event (setup_prerequisite_missing, service:setup-cli, 2s timeout) to the Datadog logs intake, the channel AI Setup uses for its own telemetry, so we can measure how often Node is missing. It needs no login, and includes the org UUID (dd_oid) when pup has a session. - Forwards pup's session as DD_ACCESS_TOKEN. When an OAuth session exists, shell API/app keys are dropped, and a lone key is never forwarded. The help lookup uses the same environment. - `--site` passed to setup becomes pup's site for the run, so the site check, the forwarded session and AI Setup's `--site` all agree. - Adds --headless when an AI agent is detected (unless --no-agent) or stdin/stdout isn't a terminal; a person at a terminal gets AI Setup's interactive flow unless they pass --headless. - Headless runs with no session point people without an account to the sign-up page and run pup's OAuth login first. At a terminal, AI Setup's own sign-in runs, and pup offers to log in too after a successful run. - Logins started by pup setup request api_keys_read, which AI Setup needs to write the org's API key into the project. A saved session without it gets one more login; `pup auth login` is unchanged. `auth::login` now returns the session it saved, so setup continues with the right session when the login lands in another region or org. - `pup setup --help` builds a setup command from AI Setup's --help --json and renders it as clap help, so new products and flags show up without a pup release. `pup setup help` and `pup help setup` are aliases. - `pup agent schema setup` says setup passes AI Setup's output through (ending with an AI_SETUP_RESULT line in headless runs) instead of returning an API body. - pup's global flags work after `setup`, and `pup setup --read-only` is blocked, though `pup setup --help` still works in read-only mode. - PUP_SETUP_AI_SETUP_PACKAGE overrides the package spec (version, tag, tarball path or file: spec) to test unreleased AI Setup builds. Co-Authored-By: Claude Opus 5.5 --- Cargo.toml | 2 +- docs/COMMANDS.md | 1 + src/commands/agent.rs | 13 + src/commands/alias.rs | 6 +- src/commands/auth.rs | 9 +- src/commands/mod.rs | 4 + src/commands/setup.rs | 1152 ++++++++++++++++++++++++++++++++++++ src/commands/setup_help.rs | 172 ++++++ src/extensions/exec.rs | 8 +- src/extensions/mod.rs | 6 + src/main.rs | 341 ++++++++++- 11 files changed, 1703 insertions(+), 11 deletions(-) create mode 100644 src/commands/setup.rs create mode 100644 src/commands/setup_help.rs 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 {