Skip to content

Commit 586760a

Browse files
phrockerphrocker
committed
fix document retrieval
1 parent d3cb9a6 commit 586760a

File tree

7 files changed

+183
-78
lines changed

7 files changed

+183
-78
lines changed

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

Lines changed: 108 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@
66
import io.sentrius.sso.core.config.SystemOptions;
77
import io.sentrius.sso.core.controllers.BaseController;
88
import io.sentrius.sso.core.dto.AgentRegistrationDTO;
9+
import io.sentrius.sso.core.dto.agents.AgentContextRequestDTO;
910
import io.sentrius.sso.core.dto.agents.AgentTemplateDTO;
1011
import io.sentrius.sso.core.exceptions.ZtatException;
12+
import io.sentrius.sso.core.model.agents.AgentContext;
1113
import io.sentrius.sso.core.model.security.enums.ApplicationAccessEnum;
1214
import io.sentrius.sso.core.services.ATPLPolicyService;
1315
import io.sentrius.sso.core.services.ErrorOutputService;
@@ -75,7 +77,7 @@ public ResponseEntity<List<AgentTemplateDTO>> getAllTemplates(
7577
if (operatingUser == null) {
7678
return ResponseEntity.status(401).build();
7779
}
78-
80+
7981
log.info("User {} requested agent templates", operatingUser.getUsername());
8082
List<AgentTemplateDTO> templates = templateService.getAllEnabledTemplates();
8183
return ResponseEntity.ok(templates);
@@ -95,7 +97,7 @@ public ResponseEntity<List<AgentTemplateDTO>> getTemplatesByCategory(
9597
if (operatingUser == null) {
9698
return ResponseEntity.status(401).build();
9799
}
98-
100+
99101
log.info("User {} requested templates for category: {}", operatingUser.getUsername(), category);
100102
List<AgentTemplateDTO> templates = templateService.getTemplatesByCategory(category);
101103
return ResponseEntity.ok(templates);
@@ -115,7 +117,7 @@ public ResponseEntity<AgentTemplateDTO> getTemplate(
115117
if (operatingUser == null) {
116118
return ResponseEntity.status(401).build();
117119
}
118-
120+
119121
return templateService.getTemplateById(id)
120122
.map(ResponseEntity::ok)
121123
.orElse(ResponseEntity.notFound().build());
@@ -135,11 +137,11 @@ public ResponseEntity<?> createTemplate(
135137
if (operatingUser == null) {
136138
return ResponseEntity.status(401).build();
137139
}
138-
140+
139141
try {
140142
templateDTO.setCreatedBy(operatingUser.getUsername());
141143
templateDTO.setSystemTemplate(false); // User templates are never system templates
142-
144+
143145
AgentTemplateDTO created = templateService.createTemplate(templateDTO);
144146
log.info("User {} created new agent template: {}", operatingUser.getUsername(), created.getName());
145147
return ResponseEntity.ok(created);
@@ -165,7 +167,7 @@ public ResponseEntity<?> updateTemplate(
165167
if (operatingUser == null) {
166168
return ResponseEntity.status(401).build();
167169
}
168-
170+
169171
try {
170172
AgentTemplateDTO updated = templateService.updateTemplate(id, templateDTO);
171173
log.info("User {} updated agent template: {}", operatingUser.getUsername(), updated.getName());
@@ -196,7 +198,7 @@ public ResponseEntity<?> deleteTemplate(
196198
if (operatingUser == null) {
197199
return ResponseEntity.status(401).build();
198200
}
199-
201+
200202
try {
201203
templateService.deleteTemplate(id);
202204
log.info("User {} deleted agent template: {}", operatingUser.getUsername(), id);
@@ -212,7 +214,7 @@ public ResponseEntity<?> deleteTemplate(
212214
.body(Map.of("error", "Failed to delete template: " + e.getMessage()));
213215
}
214216
}
215-
217+
216218
/**
217219
* Build an AgentRegistrationDTO from a template for launcher service
218220
* This endpoint provides the template configuration in a format suitable for the agent launcher
@@ -230,11 +232,11 @@ public ResponseEntity<?> prepareLaunch(
230232
if (operatingUser == null) {
231233
return ResponseEntity.status(401).build();
232234
}
233-
235+
234236
try {
235237
AgentTemplateDTO template = templateService.getTemplateById(id)
236238
.orElseThrow(() -> new IllegalArgumentException("Template not found: " + id));
237-
239+
238240
// Build AgentRegistrationDTO with full template configuration
239241
AgentRegistrationDTO agentDto = AgentRegistrationDTO.builder()
240242
.agentName(agentName)
@@ -250,10 +252,10 @@ public ResponseEntity<?> prepareLaunch(
250252
.templateLaunchConfiguration(template.getLaunchConfiguration())
251253
.agentPolicyId(template.getTrustPolicyId() != null ? template.getTrustPolicyId() : "")
252254
.build();
253-
254-
log.info("User {} prepared agent launch from template: {} -> agent: {}",
255+
256+
log.info("User {} prepared agent launch from template: {} -> agent: {}",
255257
operatingUser.getUsername(), template.getName(), agentName);
256-
258+
257259
return ResponseEntity.ok(agentDto);
258260
} catch (IllegalArgumentException e) {
259261
return ResponseEntity.notFound().build();
@@ -263,14 +265,14 @@ public ResponseEntity<?> prepareLaunch(
263265
.body(Map.of("error", "Failed to prepare launch: " + e.getMessage()));
264266
}
265267
}
266-
268+
267269
/**
268270
* Launch an agent from a template
269271
* This endpoint creates an agent registration and triggers the launcher service automatically
270-
*
272+
*
271273
* @param id Template ID
272274
* @param agentName Name for the new agent
273-
* @param agentContextId Optional context ID for the agent
275+
* @param agentContextId Optional context ID for the agent (if not provided, will be created from template)
274276
* @return Launch response with agent details
275277
*/
276278
@PostMapping("/{id}/launch")
@@ -286,18 +288,18 @@ public ResponseEntity<?> launchFromTemplate(
286288
if (operatingUser == null) {
287289
return ResponseEntity.status(401).build();
288290
}
289-
291+
290292
try {
291293
AgentTemplateDTO template = templateService.getTemplateById(id)
292294
.orElseThrow(() -> new IllegalArgumentException("Template not found: " + id));
293-
294-
log.info("User {} launching agent '{}' from template '{}'",
295+
296+
log.info("User {} launching agent '{}' from template '{}'",
295297
operatingUser.getUsername(), agentName, template.getName());
296-
298+
297299
// Check if agent is already running
298300
try {
299301
String status = agentClientService.getAgentPodStatus(
300-
appConfig.getSentriusLauncherService(),
302+
appConfig.getSentriusLauncherService(),
301303
agentName
302304
);
303305
if ("Running".equals(status) || "Pending".equals(status)) {
@@ -311,25 +313,50 @@ public ResponseEntity<?> launchFromTemplate(
311313
} catch (Exception e) {
312314
log.debug("Agent status check failed (agent may not exist yet): {}", e.getMessage());
313315
}
314-
315-
// Build AgentRegistrationDTO with full template configuration
316+
317+
// Store agent context in database if not provided
318+
String contextId = agentContextId;
319+
if (contextId == null || contextId.isEmpty()) {
320+
log.info("Creating agent context from template for agent '{}'", agentName);
321+
322+
// Build context string with all template information
323+
StringBuilder contextBuilder = new StringBuilder();
324+
contextBuilder.append("# Agent Configuration from Template: ").append(template.getName()).append("\n\n");
325+
326+
appendSectionIfPresent(contextBuilder, "Purpose", template.getPurpose(), false);
327+
appendSectionIfPresent(contextBuilder, "Goals", template.getGoals(), false);
328+
appendSectionIfPresent(contextBuilder, "Configuration", template.getDefaultConfiguration(), true);
329+
appendSectionIfPresent(contextBuilder, "Identity", template.getIdentity(), true);
330+
appendSectionIfPresent(contextBuilder, "Guardrails", template.getGuardrails(), true);
331+
appendSectionIfPresent(contextBuilder, "Launch Configuration", template.getLaunchConfiguration(), true);
332+
appendSectionIfPresent(contextBuilder, "Trust Policy ID", template.getTrustPolicyId(), false);
333+
334+
// Create agent context with template information
335+
AgentContextRequestDTO contextRequest =
336+
AgentContextRequestDTO.builder()
337+
.name(agentName)
338+
.description("Agent context created from template: " + template.getName())
339+
.context(contextBuilder.toString())
340+
.policyId(template.getTrustPolicyId())
341+
.build();
342+
343+
AgentContext savedContext = agentContextService.create(contextRequest);
344+
345+
contextId = savedContext.getId().toString();
346+
log.info("Created agent context with ID: {} for agent '{}'", contextId, agentName);
347+
}
348+
349+
// Build AgentRegistrationDTO with context ID instead of embedded template data
316350
AgentRegistrationDTO agentDto = AgentRegistrationDTO.builder()
317351
.agentName(agentName)
318352
.agentType(template.getAgentType())
319353
.agentCallbackUrl("")
320354
.clientId(agentName) // Set clientId to match agentName for policy caching
321355
.agentTemplateId(id.toString())
322-
.agentContextId(agentContextId)
323-
.templateConfiguration(template.getDefaultConfiguration())
324-
.templateIdentity(template.getIdentity())
325-
.templatePurpose(template.getPurpose())
326-
.templateGoals(template.getGoals())
327-
.templateGuardrails(template.getGuardrails())
328-
.templateTrustPolicyId(template.getTrustPolicyId())
329-
.templateLaunchConfiguration(template.getLaunchConfiguration())
356+
.agentContextId(contextId)
330357
.agentPolicyId(template.getTrustPolicyId() != null ? template.getTrustPolicyId() : "")
331358
.build();
332-
359+
333360
// Cache the policy if it exists
334361
if (template.getTrustPolicyId() != null && !template.getTrustPolicyId().isEmpty()) {
335362
var latest = atplPolicyService.getLatestPolicyEntity(template.getTrustPolicyId());
@@ -340,44 +367,42 @@ public ResponseEntity<?> launchFromTemplate(
340367
log.warn("Policy {} not found, skipping cache", template.getTrustPolicyId());
341368
}
342369
}
343-
370+
344371
// Call the launcher service
345372
zeroTrustClientService.callAuthenticatedPostOnApi(
346373
appConfig.getSentriusLauncherService(),
347374
"agent/launcher/create",
348375
agentDto
349376
);
350-
351-
// Record the agent launch if agentContextId is provided
352-
if (agentContextId != null && !agentContextId.isEmpty()) {
353-
try {
354-
UUID contextId = UUID.fromString(agentContextId);
355-
String launchedBy = operatingUser.getUserId();
356-
String parameters = String.format(
357-
"agentType=%s,templateId=%s,policyId=%s",
358-
template.getAgentType(),
359-
id.toString(),
360-
template.getTrustPolicyId() != null ? template.getTrustPolicyId() : "none"
361-
);
362-
363-
UUID launchId = agentLaunchService.recordLaunch(
364-
agentName,
365-
contextId,
366-
launchedBy,
367-
parameters
368-
);
369-
370-
log.info("Recorded agent launch: launchId={}, contextId={}, agentName={}",
371-
launchId, contextId, agentName);
372-
} catch (IllegalArgumentException e) {
373-
log.warn("Invalid agentContextId '{}', skipping launch record: {}", agentContextId, e.getMessage());
374-
} catch (Exception e) {
375-
log.warn("Failed to record agent launch (non-critical): {}", e.getMessage());
376-
}
377+
378+
// Record the agent launch with the context ID
379+
try {
380+
UUID contextUuid = UUID.fromString(contextId);
381+
String launchedBy = operatingUser.getUserId();
382+
String parameters = String.format(
383+
"agentType=%s,templateId=%s,policyId=%s",
384+
template.getAgentType(),
385+
id.toString(),
386+
template.getTrustPolicyId() != null ? template.getTrustPolicyId() : "none"
387+
);
388+
389+
UUID launchId = agentLaunchService.recordLaunch(
390+
agentName,
391+
contextUuid,
392+
launchedBy,
393+
parameters
394+
);
395+
396+
log.info("Recorded agent launch: launchId={}, contextId={}, agentName={}",
397+
launchId, contextUuid, agentName);
398+
} catch (IllegalArgumentException e) {
399+
log.warn("Invalid contextId '{}', skipping launch record: {}", contextId, e.getMessage());
400+
} catch (Exception e) {
401+
log.warn("Failed to record agent launch (non-critical): {}", e.getMessage());
377402
}
378-
403+
379404
log.info("Successfully launched agent '{}' from template '{}'", agentName, template.getName());
380-
405+
381406
return ResponseEntity.ok(Map.of(
382407
"status", "success",
383408
"message", "Agent launched successfully",
@@ -386,7 +411,7 @@ public ResponseEntity<?> launchFromTemplate(
386411
"templateName", template.getName(),
387412
"agentType", template.getAgentType()
388413
));
389-
414+
390415
} catch (IllegalArgumentException e) {
391416
return ResponseEntity.notFound().build();
392417
} catch (ZtatException e) {
@@ -399,4 +424,23 @@ public ResponseEntity<?> launchFromTemplate(
399424
.body(Map.of("error", "Failed to launch agent: " + e.getMessage()));
400425
}
401426
}
402-
}
427+
428+
/**
429+
* Helper method to append a section to the context string builder if content is present
430+
*
431+
* @param builder The StringBuilder to append to
432+
* @param title The section title
433+
* @param content The content to append
434+
* @param isJson Whether to wrap content in JSON code blocks
435+
*/
436+
private void appendSectionIfPresent(StringBuilder builder, String title, String content, boolean isJson) {
437+
if (content != null && !content.trim().isEmpty()) {
438+
builder.append("## ").append(title).append("\n");
439+
if (isJson) {
440+
builder.append("```json\n").append(content).append("\n```\n\n");
441+
} else {
442+
builder.append(content).append("\n\n");
443+
}
444+
}
445+
}
446+
}

‎core/src/main/java/io/sentrius/sso/core/dto/agents/AgentContextDTO.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,13 @@
77
import lombok.Builder;
88
import lombok.Getter;
99
import lombok.NoArgsConstructor;
10+
import lombok.ToString;
1011

1112
@Builder
1213
@Getter
1314
@NoArgsConstructor
1415
@AllArgsConstructor
15-
16+
@ToString
1617
@JsonIgnoreProperties(ignoreUnknown = true)
1718
public class AgentContextDTO {
1819
@Builder.Default

‎core/src/main/java/io/sentrius/sso/core/dto/agents/AgentContextRequestDTO.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ public class AgentContextRequestDTO {
1818
private String name;
1919
private String description;
2020
private String context;
21+
private String policyId;
2122

2223
}
2324

‎dataplane/src/main/java/io/sentrius/sso/core/model/agents/AgentContext.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,13 @@
1313
import jakarta.persistence.Table;
1414
import lombok.Getter;
1515
import lombok.Setter;
16+
import lombok.ToString;
1617

1718
@Entity
1819
@Table(name = "agent_contexts")
1920
@Getter
2021
@Setter
22+
@ToString
2123
public class AgentContext {
2224

2325
@Id

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,13 @@ public AgentContext create(@NonNull AgentContextRequestDTO dto) {
6767
context.setName(dto.getName());
6868
context.setDescription(dto.getDescription());
6969
context.setContext(refinedContext);
70+
71+
// Set policy ID if provided (handles null, empty string, and whitespace-only strings)
72+
String policyId = dto.getPolicyId();
73+
if (policyId != null && !policyId.trim().isEmpty()) {
74+
context.setPolicyId(policyId);
75+
}
76+
7077
return contextRepo.save(context);
7178
}
7279

0 commit comments

Comments
 (0)