feat: restore agent capabilities from heartbeat after server restart
The heartbeat now carries capabilities (per protocol v2 update). On each heartbeat, capabilities are updated in the agent registry. On auto-heal (server restart), capabilities from the heartbeat are used instead of empty Map.of(), so the agent's feature flags (replay, routeControl, logForwarding, etc.) are restored immediately on the first heartbeat. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -201,15 +201,17 @@ public class AgentRegistrationController {
|
||||
public ResponseEntity<Void> heartbeat(@PathVariable String id,
|
||||
@RequestBody(required = false) HeartbeatRequest request,
|
||||
HttpServletRequest httpRequest) {
|
||||
boolean found = registryService.heartbeat(id);
|
||||
Map<String, Object> capabilities = request != null ? request.getCapabilities() : null;
|
||||
boolean found = registryService.heartbeat(id, capabilities);
|
||||
if (!found) {
|
||||
// Auto-heal: re-register agent from JWT claims after server restart
|
||||
var jwtResult = (JwtService.JwtValidationResult) httpRequest.getAttribute(
|
||||
JwtAuthenticationFilter.JWT_RESULT_ATTR);
|
||||
if (jwtResult != null) {
|
||||
String application = jwtResult.application() != null ? jwtResult.application() : "default";
|
||||
Map<String, Object> caps = capabilities != null ? capabilities : Map.of();
|
||||
registryService.register(id, id, application, "unknown",
|
||||
List.of(), Map.of());
|
||||
List.of(), caps);
|
||||
registryService.heartbeat(id);
|
||||
log.info("Auto-registered agent {} (app={}) from heartbeat after server restart", id, application);
|
||||
} else {
|
||||
|
||||
@@ -55,6 +55,11 @@ public record AgentInfo(
|
||||
state, registeredAt, lastHeartbeat, newStaleTransitionTime);
|
||||
}
|
||||
|
||||
public AgentInfo withCapabilities(Map<String, Object> newCapabilities) {
|
||||
return new AgentInfo(instanceId, displayName, applicationId, version, routeIds, newCapabilities,
|
||||
state, registeredAt, lastHeartbeat, staleTransitionTime);
|
||||
}
|
||||
|
||||
public AgentInfo withMetadata(String displayName, String applicationId, String version,
|
||||
List<String> routeIds, Map<String, Object> capabilities) {
|
||||
return new AgentInfo(instanceId, displayName, applicationId, version, routeIds, capabilities,
|
||||
|
||||
@@ -71,14 +71,17 @@ public class AgentRegistryService {
|
||||
|
||||
/**
|
||||
* Process a heartbeat from an agent.
|
||||
* Updates lastHeartbeat and transitions STALE agents back to LIVE.
|
||||
* Updates lastHeartbeat, capabilities (if provided), and transitions STALE agents back to LIVE.
|
||||
*
|
||||
* @return true if the agent is known, false otherwise
|
||||
*/
|
||||
public boolean heartbeat(String id) {
|
||||
public boolean heartbeat(String id, Map<String, Object> capabilities) {
|
||||
AgentInfo updated = agents.computeIfPresent(id, (key, existing) -> {
|
||||
Instant now = Instant.now();
|
||||
AgentInfo result = existing.withLastHeartbeat(now);
|
||||
if (capabilities != null && !capabilities.isEmpty()) {
|
||||
result = result.withCapabilities(Map.copyOf(capabilities));
|
||||
}
|
||||
if (existing.state() == AgentState.STALE) {
|
||||
result = result.withState(AgentState.LIVE).withStaleTransitionTime(null);
|
||||
log.info("Agent {} revived from STALE to LIVE via heartbeat", id);
|
||||
@@ -88,6 +91,11 @@ public class AgentRegistryService {
|
||||
return updated != null;
|
||||
}
|
||||
|
||||
/** Overload for callers without capabilities (backward compatibility). */
|
||||
public boolean heartbeat(String id) {
|
||||
return heartbeat(id, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Manually transition an agent to a new state.
|
||||
* Sets staleTransitionTime when transitioning to STALE.
|
||||
|
||||
Reference in New Issue
Block a user