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()),
|
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);
|
||||||
|
|||||||
@@ -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);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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) {
|
||||||
|
|||||||
@@ -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) {
|
||||||
|
|||||||
Reference in New Issue
Block a user