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()),
AuditResult.SUCCESS, httpRequest);
// Issue JWT tokens with AGENT role
// Issue JWT tokens with AGENT role + environment
List<String> roles = List.of("AGENT");
String accessToken = jwtService.createAccessToken(request.instanceId(), application, roles);
String refreshToken = jwtService.createRefreshToken(request.instanceId(), application, roles);
String accessToken = jwtService.createAccessToken(request.instanceId(), application, environmentId, roles);
String refreshToken = jwtService.createRefreshToken(request.instanceId(), application, environmentId, roles);
return ResponseEntity.ok(new AgentRegistrationResponse(
agent.instanceId(),
@@ -181,14 +181,16 @@ public class AgentRegistrationController {
? List.of("AGENT") : result.roles();
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);
if (agent != null) {
application = agent.applicationId();
environment = agent.environmentId();
}
String newAccessToken = jwtService.createAccessToken(agentId, application, roles);
String newRefreshToken = jwtService.createRefreshToken(agentId, application, roles);
String newAccessToken = jwtService.createAccessToken(agentId, application, environment, roles);
String newRefreshToken = jwtService.createRefreshToken(agentId, application, environment, roles);
auditService.log(agentId, "agent_token_refresh", AuditCategory.AUTH, agentId,
null, AuditResult.SUCCESS, httpRequest);
@@ -211,8 +213,9 @@ public class AgentRegistrationController {
JwtAuthenticationFilter.JWT_RESULT_ATTR);
if (jwtResult != null) {
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();
registryService.register(id, id, application, "default", "unknown",
registryService.register(id, id, application, env, "unknown",
List.of(), caps);
registryService.heartbeat(id);
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);
if (jwtResult != null) {
String application = jwtResult.application() != null ? jwtResult.application() : "default";
registryService.register(id, id, application, "default", "unknown", List.of(), Map.of());
log.info("Auto-registered agent {} (app={}) from SSE connect after server restart", id, application);
String env = jwtResult.environment() != null ? jwtResult.environment() : "default";
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 {
throw new ResponseStatusException(HttpStatus.NOT_FOUND, "Agent not found: " + id);
}

View File

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

View File

@@ -18,17 +18,17 @@ public interface JwtService {
* @param application the {@code group} claim (application name)
* @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.
*/
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.
*/
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.
@@ -46,12 +46,20 @@ public interface JwtService {
// --- 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) {
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) {
return createRefreshToken(subject, application, List.of());
return createRefreshToken(subject, application, "default", List.of());
}
default String validateAndExtractAgentId(String token) {