Prepare am CLI 0.2.2 with Connected Local identity fixes

Cargo.toml

@@ -16,7 +16,7 @@ lto = "thin"
 codegen-units = 1
 
 [workspace.package]
-version = "0.2.1"
+version = "0.2.2"
 edition = "2024"
 rust-version = "1.88"
 license = "Apache-2.0"
@@ -63,7 +63,7 @@ indicatif = "0.18"
 wiremock = "0.6.5"
 clap = { version = "4", features = ["derive", "env"] }
 fs4 = { version = "0.12.0", features = ["sync"] }
-am-core-types = { path = "crates/core-types", version = "0.2.1" }
-am-cloud-types = { path = "crates/cloud-types", version = "0.2.1" }
-am-cloud-client = { path = "crates/cloud-client", version = "0.2.1" }
-atomicmemory = { path = "crates/cli", version = "0.2.1" }
+am-core-types = { path = "crates/core-types", version = "0.2.2" }
+am-cloud-types = { path = "crates/cloud-types", version = "0.2.2" }
+am-cloud-client = { path = "crates/cloud-client", version = "0.2.2" }
+atomicmemory = { path = "crates/cli", version = "0.2.2" }

crates/cli/src/commands/cloud_api_key.rs

@@ -12,7 +12,9 @@ use url::Url;
 
 use crate::auth::origin::same_origin;
 use crate::cli::GlobalOptions;
-use crate::commands::client::{cloud_api_key_client, dashboard_client};
+use crate::commands::client::{
+    cloud_api_key_client, dashboard_client, local_token_request_for_profile,
+};
 use crate::config::{
     ResolvedProfile, is_cloud_api_key, machine_scoped_key_name, require_api_key,
     require_project_id, store_api_key,
@@ -181,7 +183,7 @@ pub async fn ensure_connected_local_cloud_api_key(
         && is_cloud_api_key(&key)
     {
         if stored_key_probe_allows_reuse(
-            probe_cloud_api_key_mint(&profile.base_url, &key).await,
+            probe_cloud_api_key_mint(profile, &key).await,
             profile.project_id.as_deref(),
             &project_id,
         )? {
@@ -211,7 +213,7 @@ pub async fn ensure_connected_local_cloud_api_key(
         |secret| store_api_key(&profile.name, secret, &profile.base_url, &project_id),
     )
     .await?;
-    probe_cloud_api_key_mint(&profile.base_url, &secret)
+    probe_cloud_api_key_mint(profile, &secret)
         .await
         .context("verify newly provisioned Cloud API key (POST /v1/local/token)")?;
     if let Some(msg) = outcome.operator_message() {
@@ -228,9 +230,12 @@ pub async fn ensure_connected_local_cloud_api_key_stored(
     project_id: &str,
 ) -> Result<ProvisionOutcome> {
     if let Ok((resolved, client)) = cloud_api_key_client(global).await {
+        // Identity failures must not fall through to rotate/create — that would
+        // invalidate a still-valid stored secret without ever probing mint.
+        let req = local_token_request_for_stored_key_probe(&resolved).await?;
         // The active profile can differ from the project being provisioned.
         if stored_key_probe_allows_reuse(
-            client.mint_local_token().await.map(|_| ()),
+            client.mint_local_token(&req).await.map(|_| ()),
             resolved.project_id.as_deref(),
             project_id,
         )? {
@@ -255,6 +260,18 @@ pub async fn ensure_connected_local_cloud_api_key_stored(
     Ok(outcome)
 }
 
+/// Resolve mint identity for the stored-key reuse probe.
+///
+/// Failures preserve the stored key: callers must `?` this before create/rotate.
+async fn local_token_request_for_stored_key_probe(
+    profile: &ResolvedProfile,
+) -> Result<am_cloud_types::LocalTokenRequest> {
+    local_token_request_for_profile(profile).await.context(
+        "resolve memory_user_id for stored-key probe: stored key preserved; \
+         identity lookup failed, so setup cannot continue",
+    )
+}
+
 async fn rotate_or_create_runtime_key<F>(
     client: &dyn ConnectedLocalCredentialBackend,
     project_id: &str,
@@ -357,10 +374,16 @@ fn cloud_key_create_context(err: &CloudClientError, key_name: &str, api_origin:
     }
 }
 
-async fn probe_cloud_api_key_mint(base_url: &str, api_key: &str) -> Result<(), CloudClientError> {
-    let base = Url::parse(base_url)?;
+async fn probe_cloud_api_key_mint(
+    profile: &ResolvedProfile,
+    api_key: &str,
+) -> Result<(), CloudClientError> {
+    let base = Url::parse(&profile.base_url)?;
     let client = MemoryClient::new(base, api_key)?;
-    client.mint_local_token().await.map(|_| ())
+    let req = local_token_request_for_profile(profile)
+        .await
+        .map_err(|err| CloudClientError::Validation(err.to_string()))?;
+    client.mint_local_token(&req).await.map(|_| ())
 }
 
 #[cfg(test)]
@@ -370,6 +393,8 @@ mod tests {
 
     use super::*;
     use am_cloud_types::ApiKeyWithSecret;
+    use base64::Engine;
+    use base64::engine::general_purpose::URL_SAFE_NO_PAD;
     use chrono::{TimeZone, Utc};
 
     const TEST_LOCAL_KEY_NAME: &str = "connected-local-runtime-a1b2c3d4e5f6";
@@ -726,14 +751,21 @@ mod tests {
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
         let address = listener.local_addr().unwrap();
         let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+        let payload = URL_SAFE_NO_PAD.encode(br#"{"sub":"user_probe_test"}"#);
         let profile = ResolvedProfile {
             name: "probe-test".into(),
             base_url: format!("http://{address}"),
             kind: crate::config::ProfileKind::Cloud,
             project_id: Some("proj_test".into()),
             memory_base_url: format!("http://{address}"),
             api_key: Some("amc_stored_secret".into()),
-            oauth: None,
+            oauth: Some(crate::config::OAuthTokens {
+                id_token: format!("hdr.{payload}.sig"),
+                refresh_token: None,
+                expires_at: None,
+                issuer: None,
+                api_origin: Some(format!("http://{address}")),
+            }),
         };
         let result =
             ensure_connected_local_cloud_api_key(&GlobalOptions::default(), &profile).await;
@@ -803,3 +835,7 @@ mod tests {
         assert!(!is_api_key_quota_exceeded(&other));
     }
 }
+
+#[cfg(all(test, unix))]
+#[path = "cloud_api_key_identity_tests.rs"]
+mod identity_tests;

crates/cli/src/commands/cloud_api_key_identity_tests.rs

@@ -0,0 +1,105 @@
+//! Exercise the real stored-key entry point in a process with isolated credentials.
+
+#[path = "../../tests/support/local_token.rs"]
+mod support;
+
+use super::{ProvisionOutcome, ensure_connected_local_cloud_api_key_stored};
+use crate::cli::GlobalOptions;
+use support::Fixture;
+
+const CHILD: &str = "AM_TEST_STORED_KEY_CHILD";
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn identity_lookup_failure_preserves_valid_key() {
+    if std::env::var_os(CHILD).is_some() {
+        let result = ensure_connected_local_cloud_api_key_stored(
+            &GlobalOptions::default(),
+            "local",
+            "proj_test",
+        )
+        .await;
+        let error = result.expect_err("identity failure must not provision a key");
+        assert!(format!("{error:#}").contains("stored key preserved"));
+        return;
+    }
+    let fixture = Fixture::new().await;
+    fixture.edit("config.toml", |value| {
+        value["profiles"]["local"]
+            .as_table_mut()
+            .unwrap()
+            .remove("oauth_ref");
+    });
+    fixture.edit("credentials.toml", |value| {
+        value["oauth"]["cloud"]["expires_at"] = 0.into();
+    });
+    fixture.api.lock().unwrap().fail_first_discovery = true;
+    let before = std::fs::read(fixture.config.join("credentials.toml")).unwrap();
+    let output = fixture
+        .command(&std::env::current_exe().unwrap())
+        .env(CHILD, "1")
+        .args([
+            "--exact",
+            "commands::cloud_api_key::identity_tests::identity_lookup_failure_preserves_valid_key",
+            "--nocapture",
+        ])
+        .output()
+        .unwrap();
+    assert!(
+        output.status.success(),
+        "{}\n{}",
+        String::from_utf8_lossy(&output.stdout),
+        String::from_utf8_lossy(&output.stderr)
+    );
+    assert_eq!(
+        fixture.requests().len(),
+        1,
+        "must not retry auth then rotate an unprobed key"
+    );
+    assert_eq!(
+        fixture.requests()[0].path,
+        "/.well-known/oauth-authorization-server"
+    );
+    assert_eq!(
+        before,
+        std::fs::read(fixture.config.join("credentials.toml")).unwrap()
+    );
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn successful_mint_reuses_key_without_mutation() {
+    if std::env::var_os(CHILD).is_some() {
+        let outcome = ensure_connected_local_cloud_api_key_stored(
+            &GlobalOptions::default(),
+            "local",
+            "proj_test",
+        )
+        .await
+        .unwrap();
+        assert_eq!(outcome, ProvisionOutcome::Reused);
+        return;
+    }
+    let fixture = Fixture::new().await;
+    let before = std::fs::read(fixture.config.join("credentials.toml")).unwrap();
+    let output = fixture
+        .command(&std::env::current_exe().unwrap())
+        .env(CHILD, "1")
+        .args([
+            "--exact",
+            "commands::cloud_api_key::identity_tests::successful_mint_reuses_key_without_mutation",
+            "--nocapture",
+        ])
+        .output()
+        .unwrap();
+    assert!(
+        output.status.success(),
+        "{}\n{}",
+        String::from_utf8_lossy(&output.stdout),
+        String::from_utf8_lossy(&output.stderr)
+    );
+    assert_eq!(fixture.requests().len(), 1);
+    assert_eq!(fixture.requests()[0].path, "/v1/local/token");
+    assert_eq!(
+        before,
+        std::fs::read(fixture.config.join("credentials.toml")).unwrap()
+    );
+}

crates/cli/src/commands/connect.rs

@@ -10,7 +10,10 @@ use serde::Serialize;
 
 use crate::auth::token::valid_bearer_token;
 use crate::cli::GlobalOptions;
-use crate::commands::client::{cloud_api_key_client, dashboard_client, memory_client, resolve_ctx};
+use crate::commands::client::{
+    cloud_api_key_client, dashboard_client, local_token_request_for_profile, memory_client,
+    resolve_ctx,
+};
 use crate::commands::connect_project::{ConnectProjectOptions, run as run_connect_project};
 use crate::commands::local_clients::{
     KeyProvenance, redact_secret, render_client_env_block, resolve_local_clients,
@@ -300,7 +303,8 @@ async fn run_token(global: &GlobalOptions, print_token: bool) -> Result<()> {
     let profile = resolve_ctx(global).await?;
     ensure_local_profile(&profile)?;
     let (_profile, client) = cloud_api_key_client(global).await?;
-    let token = client.mint_local_token().await?;
+    let req = local_token_request_for_profile(&profile).await?;
+    let token = client.mint_local_token(&req).await?;
     eprintln!("warning: token printed to stdout; avoid logging or piping to files");
     println!("{}", token.access_token);
     Ok(())
@@ -483,28 +487,44 @@ async fn check_jwks_reachable(cloud_base_url: &str) -> DoctorCheck {
 
 async fn check_mint_token(global: &GlobalOptions) -> DoctorCheck {
     match cloud_api_key_client(global).await {
-        Ok((_p, client)) => match client.mint_local_token().await {
-            Ok(token) if !token.access_token.is_empty() => DoctorCheck {
-                name: "jwt_mint".into(),
-                status: "pass".into(),
-                message: format!("Cloud minted Core JWT (expires_in={}s)", token.expires_in),
-                hint: None,
-            },
-            Ok(_) => DoctorCheck {
-                name: "jwt_mint".into(),
-                status: "fail".into(),
-                message: "mint returned empty access_token".into(),
-                hint: None,
-            },
-            Err(err) => DoctorCheck {
-                name: "jwt_mint".into(),
-                status: "fail".into(),
-                message: err.to_string(),
-                hint: Some(
-                    "ensure project is type=local and Cloud API key belongs to that project".into(),
-                ),
-            },
-        },
+        Ok((profile, client)) => {
+            let req = match local_token_request_for_profile(&profile).await {
+                Ok(req) => req,
+                Err(err) => {
+                    return DoctorCheck {
+                        name: "jwt_mint".into(),
+                        status: "fail".into(),
+                        message: err.to_string(),
+                        hint: Some(
+                            "run `am auth login` so Connected Local can bind memory_user_id".into(),
+                        ),
+                    };
+                }
+            };
+            match client.mint_local_token(&req).await {
+                Ok(token) if !token.access_token.is_empty() => DoctorCheck {
+                    name: "jwt_mint".into(),
+                    status: "pass".into(),
+                    message: format!("Cloud minted Core JWT (expires_in={}s)", token.expires_in),
+                    hint: None,
+                },
+                Ok(_) => DoctorCheck {
+                    name: "jwt_mint".into(),
+                    status: "fail".into(),
+                    message: "mint returned empty access_token".into(),
+                    hint: None,
+                },
+                Err(err) => DoctorCheck {
+                    name: "jwt_mint".into(),
+                    status: "fail".into(),
+                    message: err.to_string(),
+                    hint: Some(
+                        "ensure project is type=local and Cloud API key belongs to that project"
+                            .into(),
+                    ),
+                },
+            }
+        }
         Err(err) => DoctorCheck {
             name: "jwt_mint".into(),
             status: "fail".into(),