Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b

app/src/ai/agent_sdk/federate.rs

@@ -0,0 +1,128 @@
+use std::process;
+
+use anyhow::{anyhow, Result};
+use serde_json::json;
+use warp_cli::federate::{FederateCommand, IssueGcpTokenArgs, IssueTokenArgs};
+use warp_cli::{agent::OutputFormat, GlobalOptions};
+use warp_core::{features::FeatureFlag, report_error};
+use warp_managed_secrets::ManagedSecretManager;
+use warpui::{platform::TerminationMode, AppContext, SingletonEntity as _};
+
+use super::common::set_ambient_task_context_from_run_id;
+
+/// Run identity federation commands.
+pub fn run(
+    ctx: &mut AppContext,
+    global_options: GlobalOptions,
+    command: FederateCommand,
+) -> Result<()> {
+    if !FeatureFlag::OzIdentityFederation.is_enabled() {
+        return Err(anyhow::anyhow!("This feature is not enabled"));
+    }
+    match command {
+        FederateCommand::IssueToken(args) => issue_token(ctx, args, global_options.output_format),
+        FederateCommand::IssueGcpToken(args) => issue_gcp_token(ctx, args),
+    }
+}
+
+fn issue_token(
+    ctx: &mut AppContext,
+    args: IssueTokenArgs,
+    output_format: OutputFormat,
+) -> Result<()> {
+    // Set the task ID so the ambient workload token header is sent.
+    set_ambient_task_context_from_run_id(ctx, &args.run_id)?;
+
+    let duration: std::time::Duration = args.duration.into();
+    let audience = args.audience;
+    let subject_template = match args.subject_template {
+        Some(template) => vec1::Vec1::try_from_vec(template)
+            .map_err(|_| anyhow::anyhow!("--subject-template requires at least one value"))?,
+        None => vec1::vec1!["principal".to_owned()],
+    };
+
+    ManagedSecretManager::handle(ctx).update(ctx, move |manager, ctx| {
+        let future =
+            manager.issue_task_identity_token(warp_managed_secrets::client::IdentityTokenOptions {
+                audience,
+                requested_duration: duration,
+                subject_template,
+            });
+        ctx.spawn(future, move |_, result, ctx| match result {
+            Ok(token) => {
+                let token_value = token.token;
+                let expires_at = token.expires_at.to_rfc3339();
+                let issuer = token.issuer;
+                match output_format {
+                    OutputFormat::Json | OutputFormat::Ndjson => {
+                        let output = json!({
+                            "token": token_value,
+                            "expires_at": expires_at,
+                            "issuer": issuer,
+                        });
+                        let output =
+                            serde_json::to_string(&output).expect("token output should serialize");
+                        println!("{output}");
+                    }
+                    OutputFormat::Text => {
+                        println!("{token_value}");
+                    }
+                    OutputFormat::Pretty => {
+                        println!("Token: {token_value}");
+                        println!("Expires at: {expires_at}");
+                        println!("Issuer: {issuer}");
+                    }
+                }
+                ctx.terminate_app(TerminationMode::ForceTerminate, None);
+            }
+            Err(err) => {
+                super::report_fatal_error(err, ctx);
+            }
+        });
+    });
+
+    Ok(())
+}
+fn issue_gcp_token(ctx: &mut AppContext, args: IssueGcpTokenArgs) -> Result<()> {
+    // Set the task ID so the ambient workload token header is sent.
+    set_ambient_task_context_from_run_id(ctx, &args.run_id)?;
+
+    let duration: std::time::Duration = args.duration.into();
+    let audience = args.audience;
+    let token_type = args.token_type;
+    let output_file = args.output_file;
+
+    ManagedSecretManager::handle(ctx).update(ctx, move |manager, ctx| {
+        let future =
+            manager.issue_gcp_workload_identity_federation_token(audience, token_type, duration);
+        ctx.spawn(future, move |_, result, ctx| match result {
+            Ok(token) => {
+                let output =
+                    serde_json::to_string(&token).expect("gcp token output should serialize");
+
+                // If we can't cache the token, report an error but don't fail the command.
+                if let Some(output_path) = output_file {
+                    if let Err(err) = std::fs::write(&output_path, &output) {
+                        report_error!(anyhow!(err)
+                            .context(format!("Error writing GCP token to {output_path}")));
+                    }
+                }
+
+                println!("{output}");
+                ctx.terminate_app(TerminationMode::ForceTerminate, None);
+            }
+            Err(err) => {
+                // The GCP SDK requires the executable to print a JSON error to stderr
+                // and exit with a non-zero status code. Because of this, we exit
+                // directly instead of via `ctx.terminate_app` (which would print a
+                // non-JSON error).
+                let output =
+                    serde_json::to_string(&err).expect("gcp error output should serialize");
+                eprintln!("{output}");
+                process::exit(1);
+            }
+        });
+    });
+
+    Ok(())
+}