fix(ers): preserve resolved token-chain context

service/entityresolution/integration/internal/contract_tests.go

@@ -11,6 +11,7 @@ import (
 	entityresolutionV2 "github.com/opentdf/platform/protocol/go/entityresolution/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/types/known/structpb"
 )
 
 const (
@@ -90,11 +91,12 @@ type EntityValidationRule struct {
 
 // EntityChainValidationRule defines how to validate a returned entity chain
 type EntityChainValidationRule struct {
-	EphemeralID               string   // Expected ephemeral ID
-	EntityCount               int      // Expected number of entities in the chain
-	EntityTypes               []string // Expected entity types in order
-	EntityCategories          []string // Expected entity categories in order (CATEGORY_ENVIRONMENT, CATEGORY_SUBJECT)
-	RequireConsistentOrdering bool     // Whether entity order must be consistent across implementations
+	EphemeralID               string                   // Expected ephemeral ID
+	EntityCount               int                      // Expected number of entities in the chain
+	EntityTypes               []string                 // Expected entity types in order
+	EntityCategories          []string                 // Expected entity categories in order (CATEGORY_ENVIRONMENT, CATEGORY_SUBJECT)
+	EntityRequiredFields      []map[string]interface{} // Required claims fields for each entity, by index
+	RequireConsistentOrdering bool                     // Whether entity order must be consistent across implementations
 }
 
 // ContractTestSuite holds all the contract tests for ERS implementations
@@ -521,9 +523,10 @@ func (suite *ContractTestSuite) validateContractChain(t *testing.T, chains []*en
 	entities := matchingChain.GetEntities()
 	assert.Len(t, entities, validationRule.EntityCount, "Unexpected number of entities in chain")
 
-	// Validate entity types and categories
+	// Validate entity types, categories, and preserved claims context.
 	suite.validateChainEntityTypes(t, entities, validationRule)
 	suite.validateChainEntityCategories(t, entities, validationRule)
+	suite.validateChainEntityRequiredFields(t, entities, validationRule)
 }
 
 // validateChainEntityTypes validates entity types in chain
@@ -555,6 +558,30 @@ func (suite *ContractTestSuite) validateFlexibleEntityType(t *testing.T, entitie
 	assert.Fail(t, fmt.Sprintf("Expected entity type %s not found in chain", expectedType))
 }
 
+func (suite *ContractTestSuite) validateChainEntityRequiredFields(t *testing.T, entities []*entity.Entity, validationRule EntityChainValidationRule) {
+	for idx, requiredFields := range validationRule.EntityRequiredFields {
+		if idx >= len(entities) {
+			t.Errorf("Required-fields rule index %d out of bounds (got %d entities)", idx, len(entities))
+			continue
+		}
+		if len(requiredFields) == 0 {
+			continue
+		}
+
+		claims := entities[idx].GetClaims()
+		if claims == nil {
+			t.Errorf("Entity at index %d does not contain claims", idx)
+			continue
+		}
+		var claimsStruct structpb.Struct
+		if err := claims.UnmarshalTo(&claimsStruct); err != nil {
+			t.Errorf("Failed to unpack claims for entity at index %d: %v", idx, err)
+			continue
+		}
+		suite.validateRequiredFields(t, claimsStruct.AsMap(), requiredFields)
+	}
+}
+
 // validateChainEntityCategories validates entity categories in chain
 func (suite *ContractTestSuite) validateChainEntityCategories(t *testing.T, entities []*entity.Entity, validationRule EntityChainValidationRule) {
 	for i, expectedCategory := range validationRule.EntityCategories {

service/entityresolution/integration/multistrategy_comprehensive_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/opentdf/platform/service/logger"
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/wait"
+	"google.golang.org/protobuf/types/known/structpb"
 
 	_ "github.com/lib/pq" // PostgreSQL driver
 	_ "github.com/mattn/go-sqlite3"
@@ -102,8 +103,12 @@ func TestMultiStrategy_ClaimsOnly(t *testing.T) {
 	}
 
 	entity := chain.GetEntities()[0]
-	if entity.GetUserName() != "testuser" {
-		t.Errorf("Expected username 'testuser', got '%s'", entity.GetUserName())
+	claims := chainEntityClaimsMap(t, entity)
+	if got := claims["username"]; got != "testuser" {
+		t.Errorf("Expected username 'testuser', got '%v'", got)
+	}
+	if got := claims["email_address"]; got != "test@example.com" {
+		t.Errorf("Expected email_address 'test@example.com', got '%v'", got)
 	}
 
 	t.Logf("✅ Claims-only multi-strategy test passed: Created %d entities", len(chain.GetEntities()))
@@ -267,6 +272,14 @@ func TestMultiStrategy_SQLOnly(t *testing.T) {
 		t.Fatal("Expected at least one entity in chain")
 	}
 
+	claims := chainEntityClaimsMap(t, chain.GetEntities()[0])
+	if got := claims["username"]; got != "alice" {
+		t.Fatalf("Expected username 'alice', got %v", got)
+	}
+	if got := claims["display_name"]; got != "Alice Test" {
+		t.Fatalf("Expected display_name 'Alice Test', got %v", got)
+	}
+
 	t.Logf("✅ SQL-only multi-strategy test passed: Created %d entities", len(chain.GetEntities()))
 }
 
@@ -674,8 +687,9 @@ func TestMultiStrategy_MultiProviderFailover(t *testing.T) {
 	}
 
 	entity := chain.GetEntities()[0]
-	if entity.GetUserName() != "failover-user" {
-		t.Errorf("Expected username 'failover-user', got '%s'", entity.GetUserName())
+	claims := chainEntityClaimsMap(t, entity)
+	if got := claims["username"]; got != "failover-user" {
+		t.Errorf("Expected username 'failover-user', got '%v'", got)
 	}
 
 	t.Logf("✅ Multi-provider failover test passed: Failed over to claims provider and created %d entities", len(chain.GetEntities()))
@@ -775,8 +789,9 @@ func TestMultiStrategy_MultiProviderEarlySuccess(t *testing.T) {
 	}
 
 	entity := chain.GetEntities()[0]
-	if entity.GetUserName() != "early-user" {
-		t.Errorf("Expected username 'early-user', got '%s'", entity.GetUserName())
+	claims := chainEntityClaimsMap(t, entity)
+	if got := claims["username"]; got != "early-user" {
+		t.Errorf("Expected username 'early-user', got '%v'", got)
 	}
 
 	// Should be fast because it short-circuited on first success
@@ -888,13 +903,14 @@ func TestMultiStrategy_EntityChainCreation(t *testing.T) {
 		totalEntities += len(chain.GetEntities())
 
 		entity := chain.GetEntities()[0]
+		claims := chainEntityClaimsMap(t, entity)
 		expectedUsername := fmt.Sprintf("chain-user-%d", i+1)
-		if entity.GetUserName() != expectedUsername {
-			t.Errorf("Chain %d: Expected username '%s', got '%s'", i, expectedUsername, entity.GetUserName())
+		if got := claims["username"]; got != expectedUsername {
+			t.Errorf("Chain %d: Expected username '%s', got '%v'", i, expectedUsername, got)
 		}
 
-		t.Logf("Chain %d: EphemeralId=%s, Username=%s, Entities=%d",
-			i, chain.GetEphemeralId(), entity.GetUserName(), len(chain.GetEntities()))
+		t.Logf("Chain %d: EphemeralId=%s, Username=%v, Entities=%d",
+			i, chain.GetEphemeralId(), claims["username"], len(chain.GetEntities()))
 	}
 
 	t.Logf("✅ Entity chain creation test passed: Created %d chains with %d total entities",
@@ -903,6 +919,20 @@ func TestMultiStrategy_EntityChainCreation(t *testing.T) {
 
 // Helper functions
 
+func chainEntityClaimsMap(t *testing.T, ent *entity.Entity) map[string]interface{} {
+	t.Helper()
+	claims := ent.GetClaims()
+	if claims == nil {
+		t.Fatalf("chain entity %q carries no claims; entity type is %T", ent.GetEphemeralId(), ent.GetEntityType())
+	}
+
+	var claimsStruct structpb.Struct
+	if err := claims.UnmarshalTo(&claimsStruct); err != nil {
+		t.Fatalf("failed to decode chain entity claims: %v", err)
+	}
+	return claimsStruct.AsMap()
+}
+
 func createComprehensiveTestJWT(sub, email string) string {
 	// This creates a properly formatted JWT for testing purposes (not cryptographically signed)
 	// Header: {"alg":"HS256","typ":"JWT"}
@@ -992,7 +1022,16 @@ func startSeededLDAPContainer(ctx context.Context, t *testing.T) (testcontainers
 				FileMode:          0o644,
 			},
 		},
-		WaitingFor: wait.ForListeningPort("389/tcp").WithStartupTimeout(60 * time.Second),
+		WaitingFor: wait.ForAll(
+			wait.ForListeningPort("389/tcp"),
+			wait.ForExec([]string{
+				"sh", "-c",
+				"ldapsearch -x -H ldap://localhost:389 " +
+					"-D cn=admin,dc=opentdf,dc=test -w admin123 " +
+					"-b ou=users,dc=opentdf,dc=test '(uid=alice)' dn " +
+					"| grep -q '^dn: uid=alice,ou=users,dc=opentdf,dc=test$'",
+			}),
+		).WithDeadline(60 * time.Second),
 	}
 
 	ldapContainer, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{

service/entityresolution/integration/multistrategy_test.go

@@ -445,9 +445,16 @@ func TestMultiStrategyEntityResolutionV2(t *testing.T) {
 		Expected: internal.ContractExpected{
 			ChainValidation: []internal.EntityChainValidationRule{
 				{
-					EphemeralID: "test-token-1",
-					EntityCount: 3,                                            // Multi-strategy with FailureStrategyContinue creates multiple entities from all matching strategies
-					EntityTypes: []string{"username", "username", "username"}, // All strategies create username entities
+					EphemeralID:      "test-token-1",
+					EntityCount:      3,
+					EntityTypes:      []string{"claims", "claims", "claims"},
+					EntityCategories: []string{"CATEGORY_SUBJECT", "CATEGORY_SUBJECT", "CATEGORY_SUBJECT"},
+					EntityRequiredFields: []map[string]interface{}{
+						{"username": "user123", "email": "user@example.com"},
+						{"client_id": "external-client"},
+						{"username": "user123", "email": "user@example.com"},
+					},
+					RequireConsistentOrdering: true,
 				},
 			},
 		},

service/entityresolution/integration/multistrategy_v2_serialization_test.go

@@ -23,19 +23,10 @@ import (
 // These tests exercise the real ERSV2 handler (no HTTP layer) with real
 // providers, mirroring the JustInTimePDP call patterns from production.
 
-// serializationTestConfig builds a config with a claims provider and
-// strategies covering the shapes needed to exercise the two v2 call paths:
-//   - user_strategy: matches on `sub` (JWT-cased, as CreateEntityChainsFromTokens
-//     sees it on the first call).
-//   - user_mirror_strategy: matches on `userName` (proto-cased, as the default
-//     ResolveEntities path derives it from an Entity_UserName payload — the
-//     second call in the JustInTimePDP handshake).
-//   - client_strategy: environment entity matched via `azp`.
-//
-// The mirror strategy is what the spec's "Related concern" section flags as
-// commonly missing in real configs — including it here means the two-call
-// handshake test actually hits the success path on call 2, which is where
-// the []string -> structpb bug fires.
+// serializationTestConfig builds a claims-provider configuration for both
+// direct ResolveEntities calls and token-created chains. Token chains preserve
+// the phase-1 mapped output and therefore do not need a mirror strategy for a
+// second resolution pass.
 func serializationTestConfig() types.MultiStrategyConfig {
 	return types.MultiStrategyConfig{
 		FailureStrategy: types.FailureStrategyContinue,
@@ -60,19 +51,6 @@ func serializationTestConfig() types.MultiStrategyConfig {
 					{SourceClaim: "email", ClaimName: "email_address"},
 				},
 			},
-			{
-				Name:       "user_mirror_strategy",
-				Provider:   "jwt_claims",
-				EntityType: types.EntityTypeSubject,
-				Conditions: types.StrategyConditions{
-					JWTClaims: []types.JWTClaimCondition{
-						{Claim: "userName", Operator: "exists"},
-					},
-				},
-				OutputMapping: []types.OutputMapping{
-					{SourceClaim: "userName", ClaimName: "username"},
-				},
-			},
 			{
 				Name:       "client_strategy",
 				Provider:   "jwt_claims",
@@ -178,14 +156,11 @@ func TestIntegration_ResolveEntities_ReturnsPopulatedRepresentation(t *testing.T
 	}
 }
 
-// TestIntegration_TwoCallHandshake is spec integration test 2: exact
-// reproduction of the JustInTimePDP.resolveEntitiesFromToken production
-// flow. First call CreateEntityChainsFromTokens with a signed JWT that
-// matches a configured strategy; take the returned chain, feed each
-// entity into a subsequent ResolveEntities call. Every input entity must
-// produce a non-empty representation. Today this silently drops entities
-// and the caller (KAS) can't tell.
-func TestIntegration_TwoCallHandshake(t *testing.T) {
+// TestIntegration_TokenChainPreservesResolvedClaims is spec integration test 2:
+// after CreateEntityChainsFromTokens, the returned chain should already contain
+// the resolved claims needed for downstream authz. This avoids having to route
+// the chain back through ResolveEntities from a lossy identity projection.
+func TestIntegration_TokenChainPreservesResolvedClaims(t *testing.T) {
 	if testing.Short() {
 		t.Skip("Skipping multi-strategy integration tests in short mode")
 	}
@@ -198,14 +173,13 @@ func TestIntegration_TwoCallHandshake(t *testing.T) {
 	jwt := createMockJWTForUser("alice", "alice@example.com")
 
 	chainReq := connect.NewRequest(&entityresolutionV2.CreateEntityChainsFromTokensRequest{
-		Tokens: []*entity.Token{
-			{EphemeralId: "token-alice", Jwt: jwt},
-		},
+		Tokens: []*entity.Token{{EphemeralId: "token-alice", Jwt: jwt}},
 	})
 	chainResp, err := ers.CreateEntityChainsFromTokens(t.Context(), chainReq)
 	if err != nil {
 		t.Fatalf("CreateEntityChainsFromTokens returned error: %v", err)
 	}
+
 	chains := chainResp.Msg.GetEntityChains()
 	if len(chains) != 1 {
 		t.Fatalf("EntityChains length = %d, want 1", len(chains))
@@ -215,50 +189,22 @@ func TestIntegration_TwoCallHandshake(t *testing.T) {
 		t.Fatalf("chain contains no entities")
 	}
 
-	// Second call: feed each resolved chain entity into ResolveEntities,
-	// exactly as JustInTimePDP.resolveEntitiesFromToken does. Chain
-	// entities are typed (UserName / EmailAddress / ClientId), so the
-	// handler takes the default (proto-marshalled) path — the derived
-	// claimsMap contains proto-cased names like "userName", which is
-	// why the config includes a mirror strategy keyed on "userName".
-	resolveReq := connect.NewRequest(&entityresolutionV2.ResolveEntitiesRequest{
-		Entities: chainEntities,
-	})
-	// The claims provider reads from ctx. Populate it with both the
-	// original JWT claims AND the proto-cased entity claim so the
-	// mirror strategy actually resolves (otherwise the second call would
-	// fail with "no matching strategy" and never exercise the success
-	// path where the []string bug lives).
-	resolveResp, err := ers.ResolveEntities(ctxWithClaims(t, types.JWTClaims{
-		"sub":      "alice",
-		"email":    "alice@example.com",
-		"userName": "alice",
-	}), resolveReq)
-	if err != nil {
-		t.Fatalf("ResolveEntities returned error: %v", err)
-	}
+	for i, chained := range chainEntities {
+		claims := chained.GetClaims()
+		if claims == nil {
+			t.Fatalf("entity %d should preserve resolved claims in the chain, got type %T", i, chained.GetEntityType())
+		}
 
-	reps := resolveResp.Msg.GetEntityRepresentations()
-	if len(reps) != len(chainEntities) {
-		t.Fatalf("EntityRepresentations length = %d, want %d (one per chain entity — missing entries mean the handler silently dropped them via structpb.NewStruct failure)", len(reps), len(chainEntities))
-	}
-	for i, rep := range reps {
-		props := rep.GetAdditionalProps()
-		if len(props) == 0 {
-			t.Errorf("rep[%d] AdditionalProps empty; entity %q was not resolved", i, rep.GetOriginalId())
-			continue
+		var claimsStruct structpb.Struct
+		if err := claims.UnmarshalTo(&claimsStruct); err != nil {
+			t.Fatalf("entity %d claims unmarshal failed: %v", i, err)
 		}
-		// The second call should have SUCCEEDED, not returned an error
-		// struct. If we see an error field, either the mirror strategy
-		// isn't wired up or the []string bug silently dropped a real
-		// success and something else populated an error.
-		if _, hasError := props[0].GetFields()["error"]; hasError {
-			t.Errorf("rep[%d] entity %q carries error struct instead of resolved claims: %v", i, rep.GetOriginalId(), props[0].GetFields())
+		asMap := claimsStruct.AsMap()
+		if got := asMap["username"]; got != "alice" {
+			t.Fatalf("entity %d username = %v, want alice", i, got)
 		}
-		// metadata_attempted_strategies must survive serialization — this
-		// is the field that trips the []string bug.
-		if _, ok := props[0].GetFields()["metadata_attempted_strategies"]; !ok {
-			t.Errorf("rep[%d] entity %q missing metadata_attempted_strategies — likely dropped by the structpb failure", i, rep.GetOriginalId())
+		if got := asMap["email_address"]; got != "alice@example.com" {
+			t.Fatalf("entity %d email_address = %v, want alice@example.com", i, got)
 		}
 	}
 }