From 88eaaf6d3a8cd845133b76aa233e57181488351c Mon Sep 17 00:00:00 2001
From: Hanwen Cheng <heawen.cheng@gmail.com>
Date: Thu, 7 May 2026 14:54:19 +0800
Subject: [PATCH] fix(broker): emit https://aws.amazon.com/tags claim so STS
 sets PrincipalTag

---
 .../src/handlers/oidc.rs                      | 14 +++++++++++++
 .../tests/oidc_flow.rs                        | 20 +++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/crates/agentkeys-broker-server/src/handlers/oidc.rs b/crates/agentkeys-broker-server/src/handlers/oidc.rs
index db9e913..f4137b7 100644
--- a/crates/agentkeys-broker-server/src/handlers/oidc.rs
+++ b/crates/agentkeys-broker-server/src/handlers/oidc.rs
@@ -41,6 +41,7 @@ pub async fn discovery(State(state): State<SharedState>) -> impl IntoResponse {
             "agentkeys_grant_id",
             "agentkeys_operation",
             "agentkeys_user_wallet",
+            "https://aws.amazon.com/tags",
         ],
     }))
 }
@@ -103,6 +104,13 @@ pub async fn mint_oidc_jwt(
         .unwrap_or(0);
     let exp = now + state.config.oidc_jwt_ttl_seconds as i64;
 
+    // The `https://aws.amazon.com/tags` claim is what AWS STS reads to populate
+    // session tags from the JWT. AWS does NOT auto-promote arbitrary OIDC claims
+    // — the bare `agentkeys_user_wallet` claim alone produces an untagged session,
+    // and `${aws:PrincipalTag/agentkeys_user_wallet}` in bucket policies expands
+    // to empty. `transitive_tag_keys` ensures the tag persists across role chains
+    // (e.g. assumed-role → assume-role).
+    // Spec: https://docs.aws.amazon.com/IAM/latest/UserGuide/id_session-tags.html#oidc-session-tags
     let claims = json!({
         "iss": state.config.oidc_issuer,
         "sub": format!("agentkeys:agent:{}", session.wallet),
@@ -110,6 +118,12 @@ pub async fn mint_oidc_jwt(
         "iat": now,
         "exp": exp,
         "agentkeys_user_wallet": session.wallet,
+        "https://aws.amazon.com/tags": {
+            "principal_tags": {
+                "agentkeys_user_wallet": [session.wallet],
+            },
+            "transitive_tag_keys": ["agentkeys_user_wallet"],
+        },
     });
 
     let jwt = state.oidc.sign_jwt(&claims)?;
diff --git a/crates/agentkeys-broker-server/tests/oidc_flow.rs b/crates/agentkeys-broker-server/tests/oidc_flow.rs
index 25ff1b4..2edb834 100644
--- a/crates/agentkeys-broker-server/tests/oidc_flow.rs
+++ b/crates/agentkeys-broker-server/tests/oidc_flow.rs
@@ -132,6 +132,10 @@ async fn discovery_returns_aws_compatible_shape() {
         .expect("claims_supported must be an array");
     let names: Vec<&str> = claims.iter().filter_map(|v| v.as_str()).collect();
     assert!(names.contains(&"agentkeys_user_wallet"));
+    assert!(
+        names.contains(&"https://aws.amazon.com/tags"),
+        "discovery doc must advertise the AWS tags claim so AWS IAM expects it"
+    );
     assert!(names.contains(&"sub"));
     assert!(names.contains(&"exp"));
 }
@@ -201,6 +205,22 @@ async fn mint_oidc_jwt_signs_claims_for_session_wallet() {
     assert_eq!(token_data.claims["aud"], "sts.amazonaws.com");
     assert_eq!(token_data.claims["iss"], TEST_ISSUER);
 
+    // Regression guard for the silent-Stage-7-isolation-failure bug: AWS STS
+    // populates session tags ONLY from this magic-named claim, never from
+    // arbitrary top-level claims. Without it, `${aws:PrincipalTag/...}` in
+    // bucket policies expands to empty and tenant isolation is inert.
+    let aws_tags = &token_data.claims["https://aws.amazon.com/tags"];
+    assert_eq!(
+        aws_tags["principal_tags"]["agentkeys_user_wallet"][0],
+        wallet,
+        "JWT must carry agentkeys_user_wallet as a principal_tag for STS to set the session tag"
+    );
+    assert_eq!(
+        aws_tags["transitive_tag_keys"][0],
+        "agentkeys_user_wallet",
+        "agentkeys_user_wallet must be transitive so it survives role chaining"
+    );
+
     let row = state.audit.last_row().unwrap().expect("audit row missing");
     assert_eq!(row.outcome, "ok");
     assert_eq!(row.requester_wallet, wallet);