fix: persist environment in JWT claims for auto-heal recovery
All checks were successful
CI / cleanup-branch (push) Has been skipped
CI / build (push) Successful in 1m4s
CI / docker (push) Successful in 1m7s
CI / deploy (push) Successful in 45s
CI / deploy-feature (push) Has been skipped

Add 'env' claim to agent JWTs (set at registration, carried through
refresh). Auto-heal on heartbeat/SSE now reads environment from the
JWT instead of hardcoding 'default', so agents retain their correct
environment after server restart.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
hsiegeln
2026-04-04 16:12:25 +02:00
parent 346e38ee1d
commit 72ec87a3ba
4 changed files with 36 additions and 21 deletions

View File

@@ -131,10 +131,10 @@ public class AgentRegistrationController {
Map.of("application", application, "name", request.displayName()), Map.of("application", application, "name", request.displayName()),
AuditResult.SUCCESS, httpRequest); AuditResult.SUCCESS, httpRequest);
// Issue JWT tokens with AGENT role // Issue JWT tokens with AGENT role + environment
List<String> roles = List.of("AGENT"); List<String> roles = List.of("AGENT");
String accessToken = jwtService.createAccessToken(request.instanceId(), application, roles); String accessToken = jwtService.createAccessToken(request.instanceId(), application, environmentId, roles);
String refreshToken = jwtService.createRefreshToken(request.instanceId(), application, roles); String refreshToken = jwtService.createRefreshToken(request.instanceId(), application, environmentId, roles);
return ResponseEntity.ok(new AgentRegistrationResponse( return ResponseEntity.ok(new AgentRegistrationResponse(
agent.instanceId(), agent.instanceId(),
@@ -181,14 +181,16 @@ public class AgentRegistrationController {
? List.of("AGENT") : result.roles(); ? List.of("AGENT") : result.roles();
String application = result.application() != null ? result.application() : "default"; String application = result.application() != null ? result.application() : "default";
// Try to get application from registry if available (agent may not be registered after server restart) // Try to get application + environment from registry (agent may not be registered after server restart)
String environment = result.environment() != null ? result.environment() : "default";
AgentInfo agent = registryService.findById(agentId); AgentInfo agent = registryService.findById(agentId);
if (agent != null) { if (agent != null) {
application = agent.applicationId(); application = agent.applicationId();
environment = agent.environmentId();
} }
String newAccessToken = jwtService.createAccessToken(agentId, application, roles); String newAccessToken = jwtService.createAccessToken(agentId, application, environment, roles);
String newRefreshToken = jwtService.createRefreshToken(agentId, application, roles); String newRefreshToken = jwtService.createRefreshToken(agentId, application, environment, roles);
auditService.log(agentId, "agent_token_refresh", AuditCategory.AUTH, agentId, auditService.log(agentId, "agent_token_refresh", AuditCategory.AUTH, agentId,
null, AuditResult.SUCCESS, httpRequest); null, AuditResult.SUCCESS, httpRequest);
@@ -211,8 +213,9 @@ public class AgentRegistrationController {
JwtAuthenticationFilter.JWT_RESULT_ATTR); JwtAuthenticationFilter.JWT_RESULT_ATTR);
if (jwtResult != null) { if (jwtResult != null) {
String application = jwtResult.application() != null ? jwtResult.application() : "default"; String application = jwtResult.application() != null ? jwtResult.application() : "default";
String env = jwtResult.environment() != null ? jwtResult.environment() : "default";
Map<String, Object> caps = capabilities != null ? capabilities : Map.of(); Map<String, Object> caps = capabilities != null ? capabilities : Map.of();
registryService.register(id, id, application, "default", "unknown", registryService.register(id, id, application, env, "unknown",
List.of(), caps); List.of(), caps);
registryService.heartbeat(id); registryService.heartbeat(id);
log.info("Auto-registered agent {} (app={}) from heartbeat after server restart", id, application); log.info("Auto-registered agent {} (app={}) from heartbeat after server restart", id, application);

View File

@@ -67,8 +67,9 @@ public class AgentSseController {
JwtAuthenticationFilter.JWT_RESULT_ATTR); JwtAuthenticationFilter.JWT_RESULT_ATTR);
if (jwtResult != null) { if (jwtResult != null) {
String application = jwtResult.application() != null ? jwtResult.application() : "default"; String application = jwtResult.application() != null ? jwtResult.application() : "default";
registryService.register(id, id, application, "default", "unknown", List.of(), Map.of()); String env = jwtResult.environment() != null ? jwtResult.environment() : "default";
log.info("Auto-registered agent {} (app={}) from SSE connect after server restart", id, application); registryService.register(id, id, application, env, "unknown", List.of(), Map.of());
log.info("Auto-registered agent {} (app={}, env={}) from SSE connect after server restart", id, application, env);
} else { } else {
throw new ResponseStatusException(HttpStatus.NOT_FOUND, "Agent not found: " + id); throw new ResponseStatusException(HttpStatus.NOT_FOUND, "Agent not found: " + id);
} }

View File

@@ -60,13 +60,13 @@ public class JwtServiceImpl implements JwtService {
} }
@Override @Override
public String createAccessToken(String subject, String application, List<String> roles) { public String createAccessToken(String subject, String application, String environment, List<String> roles) {
return createToken(subject, application, roles, "access", properties.getAccessTokenExpiryMs()); return createToken(subject, application, environment, roles, "access", properties.getAccessTokenExpiryMs());
} }
@Override @Override
public String createRefreshToken(String subject, String application, List<String> roles) { public String createRefreshToken(String subject, String application, String environment, List<String> roles) {
return createToken(subject, application, roles, "refresh", properties.getRefreshTokenExpiryMs()); return createToken(subject, application, environment, roles, "refresh", properties.getRefreshTokenExpiryMs());
} }
@Override @Override
@@ -84,12 +84,13 @@ public class JwtServiceImpl implements JwtService {
return validateAccessToken(token).subject(); return validateAccessToken(token).subject();
} }
private String createToken(String subject, String application, List<String> roles, private String createToken(String subject, String application, String environment,
String type, long expiryMs) { List<String> roles, String type, long expiryMs) {
Instant now = Instant.now(); Instant now = Instant.now();
JWTClaimsSet claims = new JWTClaimsSet.Builder() JWTClaimsSet claims = new JWTClaimsSet.Builder()
.subject(subject) .subject(subject)
.claim("group", application) .claim("group", application)
.claim("env", environment)
.claim("type", type) .claim("type", type)
.claim("roles", roles) .claim("roles", roles)
.issueTime(Date.from(now)) .issueTime(Date.from(now))
@@ -145,7 +146,9 @@ public class JwtServiceImpl implements JwtService {
roles = List.of(); roles = List.of();
} }
return new JwtValidationResult(subject, application, roles); String environment = claims.getStringClaim("env");
return new JwtValidationResult(subject, application, environment, roles);
} catch (ParseException e) { } catch (ParseException e) {
throw new InvalidTokenException("Failed to parse JWT", e); throw new InvalidTokenException("Failed to parse JWT", e);
} catch (JOSEException e) { } catch (JOSEException e) {

View File

@@ -18,17 +18,17 @@ public interface JwtService {
* @param application the {@code group} claim (application name) * @param application the {@code group} claim (application name)
* @param roles the {@code roles} claim (e.g. {@code ["AGENT"]}, {@code ["ADMIN"]}) * @param roles the {@code roles} claim (e.g. {@code ["AGENT"]}, {@code ["ADMIN"]})
*/ */
record JwtValidationResult(String subject, String application, List<String> roles) {} record JwtValidationResult(String subject, String application, String environment, List<String> roles) {}
/** /**
* Creates a signed access JWT with the given subject, application, and roles. * Creates a signed access JWT with the given subject, application, and roles.
*/ */
String createAccessToken(String subject, String application, List<String> roles); String createAccessToken(String subject, String application, String environment, List<String> roles);
/** /**
* Creates a signed refresh JWT with the given subject, application, and roles. * Creates a signed refresh JWT with the given subject, application, and roles.
*/ */
String createRefreshToken(String subject, String application, List<String> roles); String createRefreshToken(String subject, String application, String environment, List<String> roles);
/** /**
* Validates an access token and returns the full validation result. * Validates an access token and returns the full validation result.
@@ -46,12 +46,20 @@ public interface JwtService {
// --- Backward-compatible defaults (delegate to role-aware methods) --- // --- Backward-compatible defaults (delegate to role-aware methods) ---
default String createAccessToken(String subject, String application, List<String> roles) {
return createAccessToken(subject, application, "default", roles);
}
default String createAccessToken(String subject, String application) { default String createAccessToken(String subject, String application) {
return createAccessToken(subject, application, List.of()); return createAccessToken(subject, application, "default", List.of());
}
default String createRefreshToken(String subject, String application, List<String> roles) {
return createRefreshToken(subject, application, "default", roles);
} }
default String createRefreshToken(String subject, String application) { default String createRefreshToken(String subject, String application) {
return createRefreshToken(subject, application, List.of()); return createRefreshToken(subject, application, "default", List.of());
} }
default String validateAndExtractAgentId(String token) { default String validateAndExtractAgentId(String token) {