Skip to content

Commit 6f7bdb8

Browse files
Copilotphrocker
Copilot
and
phrocker
authored
Fix lineage display, add API endpoint to create generational agents, and integrate launch tracking (#99)
* Initial plan * Fix lineage to show complete family tree (ancestors and descendants) Co-authored-by: phrocker <[email protected]> * Add API endpoint to create next generation agents Co-authored-by: phrocker <[email protected]> * Record agent launches in agent_launches table when contextId is provided Co-authored-by: phrocker <[email protected]> --------- Co-authored-by: copilot-swe-agent[bot] <[email protected]> Co-authored-by: phrocker <[email protected]>
1 parent 03593d5 commit 6f7bdb8

File tree

5 files changed

+346
-7
lines changed

5 files changed

+346
-7
lines changed

‎api/src/main/java/io/sentrius/sso/controllers/api/agents/AgentApiController.java‎

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import io.sentrius.sso.core.services.agents.AgentClientService;
4141
import io.sentrius.sso.core.services.agents.AgentContextService;
4242
import io.sentrius.sso.core.services.agents.AgentService;
43+
import io.sentrius.sso.core.services.agents.GenerationManager;
4344
import io.sentrius.sso.core.services.auditing.AuditService;
4445
import io.sentrius.sso.core.services.security.CryptoService;
4546
import io.sentrius.sso.core.services.security.KeycloakService;
@@ -55,6 +56,7 @@
5556
import lombok.extern.slf4j.Slf4j;
5657
import org.apache.http.HttpStatus;
5758
import org.jetbrains.annotations.NotNull;
59+
import org.springframework.beans.factory.annotation.Autowired;
5860
import org.springframework.data.domain.PageRequest;
5961
import org.springframework.data.domain.Pageable;
6062
import org.springframework.data.domain.Sort;
@@ -87,6 +89,7 @@ public class AgentApiController extends BaseController {
8789
final AgentContextService agentContextService;
8890
final AgentClientService agentClientService;
8991
final AppConfig appConfig;
92+
final GenerationManager generationManager;
9093

9194
public AgentApiController(
9295
UserService userService,
@@ -98,7 +101,8 @@ public AgentApiController(
98101
ZeroTrustAccessTokenService ztatService, ZeroTrustRequestService ztrService, AgentService agentService,
99102
ProvenanceKafkaProducer provenanceKafkaProducer, ZeroTrustRequestService ztatRequestService,
100103
AgentContextService agentContextService, AgentClientService agentClientService,
101-
AppConfig appConfig
104+
AppConfig appConfig,
105+
@Autowired(required = false) GenerationManager generationManager
102106
) {
103107
super(userService, systemOptions, errorOutputService);
104108
this.auditService = auditService;
@@ -114,6 +118,7 @@ public AgentApiController(
114118
this.agentContextService = agentContextService;
115119
this.agentClientService = agentClientService;
116120
this.appConfig = appConfig;
121+
this.generationManager = generationManager;
117122
}
118123

119124
public SessionLog createSession(@RequestParam String username, @RequestParam String ipAddress) {
@@ -897,6 +902,66 @@ public ResponseEntity<AgentContextDTO> createContext(
897902
return ResponseEntity.ok(dto);
898903
}
899904

905+
@PostMapping("/context/{contextId}/generation")
906+
@LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION})
907+
public ResponseEntity<?> createNextGeneration(
908+
HttpServletRequest request,
909+
HttpServletResponse response,
910+
@PathVariable("contextId") String contextId) {
911+
912+
if (generationManager == null) {
913+
log.error("GenerationManager not available");
914+
return ResponseEntity.status(HttpStatus.SC_SERVICE_UNAVAILABLE)
915+
.body(Map.of("error", "Generation management not available"));
916+
}
917+
918+
try {
919+
var operatingUser = getOperatingUser(request, response);
920+
if (operatingUser == null) {
921+
log.warn("No operating user found");
922+
return ResponseEntity.status(HttpStatus.SC_UNAUTHORIZED)
923+
.body(Map.of("error", "Unauthorized"));
924+
}
925+
926+
UUID parentId = UUID.fromString(contextId);
927+
var childContext = generationManager.createNextGeneration(parentId, operatingUser.getUserId());
928+
929+
long inheritedCount = agentContextService.getInheritedMemoryCount(childContext.getId());
930+
931+
var dto = AgentContextDTO.builder()
932+
.contextId(childContext.getId())
933+
.name(childContext.getName())
934+
.description(childContext.getDescription())
935+
.context(childContext.getContext())
936+
.createdAt(childContext.getCreatedAt())
937+
.updatedAt(childContext.getUpdatedAt())
938+
.generation(childContext.getGeneration())
939+
.parentId(childContext.getParentId())
940+
.memoryNamespace(childContext.getMemoryNamespace())
941+
.trustScore(childContext.getTrustScore())
942+
.policyId(childContext.getPolicyId())
943+
.inheritedMemoryCount(inheritedCount)
944+
.build();
945+
946+
log.info("Created next generation agent: gen={}, parent={}, child={}",
947+
childContext.getGeneration(), parentId, childContext.getId());
948+
return ResponseEntity.ok(dto);
949+
950+
} catch (IllegalArgumentException e) {
951+
log.error("Invalid parent context ID: {}", contextId, e);
952+
return ResponseEntity.status(HttpStatus.SC_NOT_FOUND)
953+
.body(Map.of("error", "Parent agent context not found: " + contextId));
954+
} catch (IllegalStateException e) {
955+
log.error("Failed to create next generation: {}", e.getMessage(), e);
956+
return ResponseEntity.status(HttpStatus.SC_FORBIDDEN)
957+
.body(Map.of("error", e.getMessage()));
958+
} catch (Exception e) {
959+
log.error("Unexpected error creating next generation", e);
960+
return ResponseEntity.status(HttpStatus.SC_INTERNAL_SERVER_ERROR)
961+
.body(Map.of("error", "Failed to create next generation: " + e.getMessage()));
962+
}
963+
}
964+
900965
@GetMapping("/stats")
901966
@LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_LOG_IN})
902967
public ResponseEntity<?> getAgentStats(HttpServletRequest request, HttpServletResponse response) throws ZtatException {

‎api/src/main/java/io/sentrius/sso/controllers/api/agents/AgentBootstrapController.java‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import java.security.GeneralSecurityException;
66
import java.util.List;
77
import java.util.Optional;
8+
import java.util.UUID;
89
import com.fasterxml.jackson.databind.ObjectMapper;
910
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
1011
import com.google.common.collect.Maps;
@@ -24,6 +25,7 @@
2425
import io.sentrius.sso.core.services.ErrorOutputService;
2526
import io.sentrius.sso.core.services.UserService;
2627
import io.sentrius.sso.core.services.agents.AgentClientService;
28+
import io.sentrius.sso.core.services.agents.AgentLaunchService;
2729
import io.sentrius.sso.core.services.agents.AgentService;
2830
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
2931
import io.sentrius.sso.core.services.auditing.AuditService;
@@ -61,6 +63,7 @@ public class AgentBootstrapController extends BaseController {
6163
private final ObjectMapper yamlMapper = new ObjectMapper(new YAMLFactory());
6264
final AppConfig appConfig;
6365
private final AgentClientService agentClientService;
66+
private final AgentLaunchService agentLaunchService;
6467

6568

6669
public AgentBootstrapController(
@@ -72,7 +75,8 @@ public AgentBootstrapController(
7275
ATPLPolicyService atplPolicyService,
7376
ZeroTrustAccessTokenService ztatService, ZeroTrustRequestService ztrService, AgentService agentService,
7477
ZeroTrustClientService zeroTrustClientService, AppConfig appConfig,
75-
AgentClientService agentClientService
78+
AgentClientService agentClientService,
79+
AgentLaunchService agentLaunchService
7680
) {
7781
super(userService, systemOptions, errorOutputService);
7882
this.auditService = auditService;
@@ -86,6 +90,7 @@ public AgentBootstrapController(
8690
this.zeroTrustClientService = zeroTrustClientService;
8791
this.appConfig = appConfig;
8892
this.agentClientService = agentClientService;
93+
this.agentLaunchService = agentLaunchService;
8994
}
9095

9196

@@ -219,6 +224,34 @@ public ResponseEntity<String> launchPod(
219224
var operatingUser = getOperatingUser(request, response );
220225
zeroTrustClientService.callAuthenticatedPostOnApi(appConfig.getSentriusLauncherService(), "agent/launcher/create",
221226
registrationDTO);
227+
228+
// Record the agent launch if agentContextId is provided
229+
if (registrationDTO.getAgentContextId() != null && !registrationDTO.getAgentContextId().isEmpty()) {
230+
try {
231+
UUID contextId = UUID.fromString(registrationDTO.getAgentContextId());
232+
String launchedBy = operatingUser != null ? operatingUser.getUserId() : "system";
233+
String parameters = "agentType=" + registrationDTO.getAgentType() +
234+
",policyId=" + registrationDTO.getAgentPolicyId();
235+
236+
UUID launchId = agentLaunchService.recordLaunch(
237+
registrationDTO.getAgentName(),
238+
contextId,
239+
launchedBy,
240+
parameters
241+
);
242+
243+
log.info("Recorded agent launch: launchId={}, contextId={}, agentName={}",
244+
launchId, contextId, registrationDTO.getAgentName());
245+
} catch (IllegalArgumentException e) {
246+
log.error("Invalid agentContextId: {}", registrationDTO.getAgentContextId(), e);
247+
// Don't fail the launch, just log the error
248+
} catch (Exception e) {
249+
log.error("Failed to record agent launch", e);
250+
// Don't fail the launch, just log the error
251+
}
252+
} else {
253+
log.info("No agentContextId provided, skipping launch record for {}", registrationDTO.getAgentName());
254+
}
222255

223256
// bootstrap with a default policy
224257
return ResponseEntity.ok("{\"status\": \"success\"}");

‎dataplane/src/main/java/io/sentrius/sso/core/repository/AgentContextRepository.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@
22

33
import io.sentrius.sso.core.model.agents.AgentContext;
44
import org.springframework.data.jpa.repository.JpaRepository;
5+
import java.util.List;
56
import java.util.Optional;
67
import java.util.UUID;
78

89
public interface AgentContextRepository extends JpaRepository<AgentContext, UUID> {
910
Optional<AgentContext> findByName(String name);
11+
List<AgentContext> findByParentId(UUID parentId);
1012
}
1113

‎dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentContextService.java‎

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,17 +45,49 @@ public List<AgentContext> getLineage(UUID agentId) {
4545
List<AgentContext> lineage = new ArrayList<>();
4646
AgentContext current = contextRepo.findById(agentId).orElse(null);
4747

48-
while (current != null) {
49-
lineage.add(0, current);
50-
if (current.getParentId() != null) {
51-
current = contextRepo.findById(current.getParentId()).orElse(null);
48+
if (current == null) {
49+
return lineage;
50+
}
51+
52+
// First, traverse up to find the root ancestor
53+
AgentContext root = current;
54+
List<AgentContext> ancestors = new ArrayList<>();
55+
while (root.getParentId() != null) {
56+
AgentContext parent = contextRepo.findById(root.getParentId()).orElse(null);
57+
if (parent != null) {
58+
ancestors.add(0, parent);
59+
root = parent;
5260
} else {
53-
current = null;
61+
break;
5462
}
5563
}
5664

65+
// Add all ancestors to lineage
66+
lineage.addAll(ancestors);
67+
68+
// Add the current agent if not already in lineage
69+
if (!lineage.contains(current)) {
70+
lineage.add(current);
71+
}
72+
73+
// Now traverse down from current to find all descendants recursively
74+
addDescendants(current, lineage);
75+
5776
return lineage;
5877
}
78+
79+
/**
80+
* Recursively adds all descendants of the given agent to the lineage list.
81+
*/
82+
private void addDescendants(AgentContext parent, List<AgentContext> lineage) {
83+
List<AgentContext> children = contextRepo.findByParentId(parent.getId());
84+
for (AgentContext child : children) {
85+
if (!lineage.contains(child)) {
86+
lineage.add(child);
87+
addDescendants(child, lineage);
88+
}
89+
}
90+
}
5991

6092
public List<AgentContext> getLineageByName(String agentName) {
6193
log.info("Getting lineage for agent by name: {}", agentName);

0 commit comments

Comments
 (0)