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,
|
public ResponseEntity<Void> heartbeat(@PathVariable String id,
|
||||||
@RequestBody(required = false) HeartbeatRequest request,
|
@RequestBody(required = false) HeartbeatRequest request,
|
||||||
HttpServletRequest httpRequest) {
|
HttpServletRequest httpRequest) {
|
||||||
boolean found = registryService.heartbeat(id);
|
Map<String, Object> capabilities = request != null ? request.getCapabilities() : null;
|
||||||
|
boolean found = registryService.heartbeat(id, capabilities);
|
||||||
if (!found) {
|
if (!found) {
|
||||||
// Auto-heal: re-register agent from JWT claims after server restart
|
// Auto-heal: re-register agent from JWT claims after server restart
|
||||||
var jwtResult = (JwtService.JwtValidationResult) httpRequest.getAttribute(
|
var jwtResult = (JwtService.JwtValidationResult) httpRequest.getAttribute(
|
||||||
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";
|
||||||
|
Map<String, Object> caps = capabilities != null ? capabilities : Map.of();
|
||||||
registryService.register(id, id, application, "unknown",
|
registryService.register(id, id, application, "unknown",
|
||||||
List.of(), Map.of());
|
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);
|
||||||
} else {
|
} else {
|
||||||
|
|||||||
@@ -55,6 +55,11 @@ public record AgentInfo(
|
|||||||
state, registeredAt, lastHeartbeat, newStaleTransitionTime);
|
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,
|
public AgentInfo withMetadata(String displayName, String applicationId, String version,
|
||||||
List<String> routeIds, Map<String, Object> capabilities) {
|
List<String> routeIds, Map<String, Object> capabilities) {
|
||||||
return new AgentInfo(instanceId, displayName, applicationId, version, routeIds, capabilities,
|
return new AgentInfo(instanceId, displayName, applicationId, version, routeIds, capabilities,
|
||||||
|
|||||||
@@ -71,14 +71,17 @@ public class AgentRegistryService {
|
|||||||
|
|
||||||
/**
|
/**
|
||||||
* Process a heartbeat from an agent.
|
* 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
|
* @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) -> {
|
AgentInfo updated = agents.computeIfPresent(id, (key, existing) -> {
|
||||||
Instant now = Instant.now();
|
Instant now = Instant.now();
|
||||||
AgentInfo result = existing.withLastHeartbeat(now);
|
AgentInfo result = existing.withLastHeartbeat(now);
|
||||||
|
if (capabilities != null && !capabilities.isEmpty()) {
|
||||||
|
result = result.withCapabilities(Map.copyOf(capabilities));
|
||||||
|
}
|
||||||
if (existing.state() == AgentState.STALE) {
|
if (existing.state() == AgentState.STALE) {
|
||||||
result = result.withState(AgentState.LIVE).withStaleTransitionTime(null);
|
result = result.withState(AgentState.LIVE).withStaleTransitionTime(null);
|
||||||
log.info("Agent {} revived from STALE to LIVE via heartbeat", id);
|
log.info("Agent {} revived from STALE to LIVE via heartbeat", id);
|
||||||
@@ -88,6 +91,11 @@ public class AgentRegistryService {
|
|||||||
return updated != null;
|
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.
|
* Manually transition an agent to a new state.
|
||||||
* Sets staleTransitionTime when transitioning to STALE.
|
* Sets staleTransitionTime when transitioning to STALE.
|
||||||
|
|||||||
Reference in New Issue
Block a user