fix: persist environment in JWT claims for auto-heal recovery
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:
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user