From 72ec87a3ba3c467ce46e2c1d16c23b55db8494e5 Mon Sep 17 00:00:00 2001 From: hsiegeln <37154749+hsiegeln@users.noreply.github.com> Date: Sat, 4 Apr 2026 16:12:25 +0200 Subject: [PATCH] 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) --- .../AgentRegistrationController.java | 17 ++++++++++------- .../app/controller/AgentSseController.java | 5 +++-- .../server/app/security/JwtServiceImpl.java | 17 ++++++++++------- .../server/core/security/JwtService.java | 18 +++++++++++++----- 4 files changed, 36 insertions(+), 21 deletions(-) diff --git a/cameleer3-server-app/src/main/java/com/cameleer3/server/app/controller/AgentRegistrationController.java b/cameleer3-server-app/src/main/java/com/cameleer3/server/app/controller/AgentRegistrationController.java index 15fcfbd5..79b0de86 100644 --- a/cameleer3-server-app/src/main/java/com/cameleer3/server/app/controller/AgentRegistrationController.java +++ b/cameleer3-server-app/src/main/java/com/cameleer3/server/app/controller/AgentRegistrationController.java @@ -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 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 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); diff --git a/cameleer3-server-app/src/main/java/com/cameleer3/server/app/controller/AgentSseController.java b/cameleer3-server-app/src/main/java/com/cameleer3/server/app/controller/AgentSseController.java index 548674cb..94dd0735 100644 --- a/cameleer3-server-app/src/main/java/com/cameleer3/server/app/controller/AgentSseController.java +++ b/cameleer3-server-app/src/main/java/com/cameleer3/server/app/controller/AgentSseController.java @@ -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); } diff --git a/cameleer3-server-app/src/main/java/com/cameleer3/server/app/security/JwtServiceImpl.java b/cameleer3-server-app/src/main/java/com/cameleer3/server/app/security/JwtServiceImpl.java index b04a16f3..ca4d9745 100644 --- a/cameleer3-server-app/src/main/java/com/cameleer3/server/app/security/JwtServiceImpl.java +++ b/cameleer3-server-app/src/main/java/com/cameleer3/server/app/security/JwtServiceImpl.java @@ -60,13 +60,13 @@ public class JwtServiceImpl implements JwtService { } @Override - public String createAccessToken(String subject, String application, List roles) { - return createToken(subject, application, roles, "access", properties.getAccessTokenExpiryMs()); + public String createAccessToken(String subject, String application, String environment, List roles) { + return createToken(subject, application, environment, roles, "access", properties.getAccessTokenExpiryMs()); } @Override - public String createRefreshToken(String subject, String application, List roles) { - return createToken(subject, application, roles, "refresh", properties.getRefreshTokenExpiryMs()); + public String createRefreshToken(String subject, String application, String environment, List 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 roles, - String type, long expiryMs) { + private String createToken(String subject, String application, String environment, + List 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) { diff --git a/cameleer3-server-core/src/main/java/com/cameleer3/server/core/security/JwtService.java b/cameleer3-server-core/src/main/java/com/cameleer3/server/core/security/JwtService.java index a10b617f..17fb73e2 100644 --- a/cameleer3-server-core/src/main/java/com/cameleer3/server/core/security/JwtService.java +++ b/cameleer3-server-core/src/main/java/com/cameleer3/server/core/security/JwtService.java @@ -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 roles) {} + record JwtValidationResult(String subject, String application, String environment, List roles) {} /** * Creates a signed access JWT with the given subject, application, and roles. */ - String createAccessToken(String subject, String application, List roles); + String createAccessToken(String subject, String application, String environment, List roles); /** * Creates a signed refresh JWT with the given subject, application, and roles. */ - String createRefreshToken(String subject, String application, List roles); + String createRefreshToken(String subject, String application, String environment, List 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 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 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) {