feat: restore agent capabilities from heartbeat after server restart
All checks were successful
CI / cleanup-branch (push) Has been skipped
CI / build (push) Successful in 1m8s
CI / docker (push) Successful in 40s
CI / deploy-feature (push) Has been skipped
CI / deploy (push) Successful in 37s

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:
hsiegeln
2026-04-04 13:19:15 +02:00
parent abed4dc96f
commit 45a74075a1
3 changed files with 19 additions and 4 deletions

View File

@@ -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 {

View File

@@ -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,

View File

@@ -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.