From 65c8beea91de9dfe1336073b7feabd530651c539 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 8 Jul 2025 00:29:02 +0000 Subject: [PATCH 1/9] Initial plan From 42fbb3ee647d4016c6180f361c0da354ccfb73b4 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 8 Jul 2025 00:48:42 +0000 Subject: [PATCH 2/9] Implement ATPL chat configuration system with enhanced capabilities model Co-authored-by: phrocker <1781585+phrocker@users.noreply.github.com> --- .../controllers/api/ATPLChatController.java | 206 +++++++ .../view/ATPLConfigController.java | 6 + .../sso/services/ATPLChatService.java | 262 +++++++++ .../resources/templates/sso/atpl/chat.html | 555 ++++++++++++++++++ .../resources/templates/sso/atpl/list.html | 9 +- .../sentrius/sso/core/trust/Capability.java | 6 + 6 files changed, 1043 insertions(+), 1 deletion(-) create mode 100644 api/src/main/java/io/sentrius/sso/controllers/api/ATPLChatController.java create mode 100644 api/src/main/java/io/sentrius/sso/services/ATPLChatService.java create mode 100644 api/src/main/resources/templates/sso/atpl/chat.html diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/ATPLChatController.java b/api/src/main/java/io/sentrius/sso/controllers/api/ATPLChatController.java new file mode 100644 index 00000000..45524e57 --- /dev/null +++ b/api/src/main/java/io/sentrius/sso/controllers/api/ATPLChatController.java @@ -0,0 +1,206 @@ +package io.sentrius.sso.controllers.api; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import io.sentrius.sso.core.annotations.LimitAccess; +import io.sentrius.sso.core.config.SystemOptions; +import io.sentrius.sso.core.controllers.BaseController; +import io.sentrius.sso.core.model.security.enums.ApplicationAccessEnum; +import io.sentrius.sso.core.services.ATPLPolicyService; +import io.sentrius.sso.core.services.ErrorOutputService; +import io.sentrius.sso.core.services.UserService; +import io.sentrius.sso.core.trust.ATPLPolicy; +import io.sentrius.sso.services.ATPLChatService; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import lombok.extern.slf4j.Slf4j; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.*; + +@Slf4j +@RestController +@RequestMapping("/api/v1/atpl/chat") +public class ATPLChatController extends BaseController { + + private final ATPLChatService atplChatService; + private final ATPLPolicyService atplPolicyService; + private final ObjectMapper objectMapper; + + // Store conversation sessions (in production, use Redis or database) + private final Map> chatSessions = new HashMap<>(); + + public ATPLChatController( + UserService userService, + SystemOptions systemOptions, + ErrorOutputService errorOutputService, + ATPLChatService atplChatService, + ATPLPolicyService atplPolicyService, + ObjectMapper objectMapper + ) { + super(userService, systemOptions, errorOutputService); + this.atplChatService = atplChatService; + this.atplPolicyService = atplPolicyService; + this.objectMapper = objectMapper; + } + + @PostMapping("/message") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public ResponseEntity processMessage( + @RequestBody Map request, + HttpServletRequest httpRequest, + HttpServletResponse httpResponse + ) { + try { + String sessionId = (String) request.get("sessionId"); + String message = (String) request.get("message"); + + if (sessionId == null || message == null) { + return ResponseEntity.badRequest().body(createErrorResponse("Session ID and message are required")); + } + + // Get or create session context + Map sessionContext = chatSessions.computeIfAbsent(sessionId, k -> new HashMap<>()); + + // Process the message with the ATPL chat service + String response = atplChatService.processATPLChatMessage(message, sessionContext); + + // Update session context with conversation history + updateSessionContext(sessionContext, message, response); + + ObjectNode result = objectMapper.createObjectNode(); + result.put("response", response); + result.put("sessionId", sessionId); + + return ResponseEntity.ok(result); + + } catch (Exception e) { + log.error("Error processing ATPL chat message: {}", e.getMessage(), e); + return ResponseEntity.internalServerError().body(createErrorResponse("Failed to process message")); + } + } + + @PostMapping("/generate-policy") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public ResponseEntity generatePolicy( + @RequestBody Map request, + HttpServletRequest httpRequest, + HttpServletResponse httpResponse + ) { + try { + String sessionId = (String) request.get("sessionId"); + Map sessionContext = chatSessions.get(sessionId); + + if (sessionContext == null) { + return ResponseEntity.badRequest().body(createErrorResponse("Session not found")); + } + + ObjectNode policyNode = atplChatService.generateATPLPolicy(sessionContext); + + ObjectNode result = objectMapper.createObjectNode(); + result.set("policy", policyNode); + result.put("sessionId", sessionId); + + return ResponseEntity.ok(result); + + } catch (Exception e) { + log.error("Error generating ATPL policy: {}", e.getMessage(), e); + return ResponseEntity.internalServerError().body(createErrorResponse("Failed to generate policy")); + } + } + + @GetMapping("/suggestions") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public ResponseEntity getSuggestions( + @RequestParam String description, + HttpServletRequest httpRequest, + HttpServletResponse httpResponse + ) { + try { + List suggestions = atplChatService.suggestCapabilities(description); + + ObjectNode result = objectMapper.createObjectNode(); + ArrayNode suggestionsArray = objectMapper.createArrayNode(); + suggestions.forEach(suggestionsArray::add); + result.set("suggestions", suggestionsArray); + + return ResponseEntity.ok(result); + + } catch (Exception e) { + log.error("Error getting suggestions: {}", e.getMessage(), e); + return ResponseEntity.internalServerError().body(createErrorResponse("Failed to get suggestions")); + } + } + + @GetMapping("/existing-policies") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public ResponseEntity getExistingPolicies( + HttpServletRequest httpRequest, + HttpServletResponse httpResponse + ) { + try { + List policies = atplPolicyService.getAllPolicies(); + + ObjectNode result = objectMapper.createObjectNode(); + ArrayNode policiesArray = objectMapper.createArrayNode(); + + for (ATPLPolicy policy : policies) { + ObjectNode policyNode = objectMapper.createObjectNode(); + policyNode.put("id", policy.getPolicyId()); + policyNode.put("description", policy.getDescription()); + policyNode.put("version", policy.getVersion()); + + if (policy.getCapabilities() != null && policy.getCapabilities().getPrimitives() != null) { + ArrayNode capabilitiesArray = objectMapper.createArrayNode(); + policy.getCapabilities().getPrimitives().forEach(cap -> { + ObjectNode capNode = objectMapper.createObjectNode(); + capNode.put("id", cap.getId()); + capNode.put("description", cap.getDescription()); + capabilitiesArray.add(capNode); + }); + policyNode.set("capabilities", capabilitiesArray); + } + + policiesArray.add(policyNode); + } + + result.set("policies", policiesArray); + return ResponseEntity.ok(result); + + } catch (Exception e) { + log.error("Error getting existing policies: {}", e.getMessage(), e); + return ResponseEntity.internalServerError().body(createErrorResponse("Failed to get existing policies")); + } + } + + private void updateSessionContext(Map sessionContext, String userMessage, String agentResponse) { + // Store conversation history + @SuppressWarnings("unchecked") + List> history = (List>) sessionContext.computeIfAbsent("conversation_history", k -> new java.util.ArrayList<>()); + + Map userEntry = new HashMap<>(); + userEntry.put("type", "user"); + userEntry.put("message", userMessage); + history.add(userEntry); + + Map agentEntry = new HashMap<>(); + agentEntry.put("type", "assistant"); + agentEntry.put("message", agentResponse); + history.add(agentEntry); + + // Keep only last 10 exchanges + if (history.size() > 20) { + history.subList(0, history.size() - 20).clear(); + } + } + + private ObjectNode createErrorResponse(String message) { + ObjectNode error = objectMapper.createObjectNode(); + error.put("error", message); + return error; + } +} \ No newline at end of file diff --git a/api/src/main/java/io/sentrius/sso/controllers/view/ATPLConfigController.java b/api/src/main/java/io/sentrius/sso/controllers/view/ATPLConfigController.java index 58b0414e..618f5241 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/view/ATPLConfigController.java +++ b/api/src/main/java/io/sentrius/sso/controllers/view/ATPLConfigController.java @@ -64,4 +64,10 @@ public String configurePage(Model model, @RequestParam(name= "id", required=fals } return "sso/atpl/configure"; } + + @GetMapping("/chat") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public String chatPage(Model model) { + return "sso/atpl/chat"; + } } \ No newline at end of file diff --git a/api/src/main/java/io/sentrius/sso/services/ATPLChatService.java b/api/src/main/java/io/sentrius/sso/services/ATPLChatService.java new file mode 100644 index 00000000..851f2cdd --- /dev/null +++ b/api/src/main/java/io/sentrius/sso/services/ATPLChatService.java @@ -0,0 +1,262 @@ +package io.sentrius.sso.services; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import io.sentrius.sso.core.services.ATPLPolicyService; +import io.sentrius.sso.core.trust.ATPLPolicy; +import io.sentrius.sso.core.trust.CapabilitySet; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Service; + +@Slf4j +@Service +@RequiredArgsConstructor +public class ATPLChatService { + + private final ATPLPolicyService atplPolicyService; + private final ObjectMapper objectMapper; + + public String processATPLChatMessage(String userMessage, Map context) { + String lowerMessage = userMessage.toLowerCase(); + + // Simple rule-based responses + if (lowerMessage.contains("start") || lowerMessage.contains("begin") || lowerMessage.contains("new")) { + return "I'll help you create a new ATPL policy. Let's start by understanding what your agent needs to do. " + + "What are the main capabilities or functions your agent should have? For example:\n" + + "• File operations (read, write, delete)\n" + + "• System monitoring (CPU, memory, disk)\n" + + "• Network communications\n" + + "• Database access\n" + + "• External API calls"; + } + + if (lowerMessage.contains("endpoint")) { + List existingPolicies = atplPolicyService.getAllPolicies(); + StringBuilder response = new StringBuilder(); + response.append("For endpoints, consider what APIs or services your agent needs to access. "); + + if (!existingPolicies.isEmpty()) { + response.append("Here are some endpoints from existing policies:\n"); + existingPolicies.forEach(policy -> { + if (policy.getCapabilities() != null && policy.getCapabilities().getPrimitives() != null) { + policy.getCapabilities().getPrimitives().forEach(cap -> { + if (cap.getEndpoints() != null && !cap.getEndpoints().isEmpty()) { + response.append("• ").append(String.join(", ", cap.getEndpoints())).append("\n"); + } + }); + } + }); + } else { + response.append("Common examples include:\n"); + response.append("• /api/v1/data/read - for reading data\n"); + response.append("• /api/v1/system/status - for system status\n"); + response.append("• /api/v1/execute - for executing operations\n"); + } + + response.append("\nWhat specific endpoints does your agent need to access?"); + return response.toString(); + } + + if (lowerMessage.contains("command")) { + return "For commands, think about what system commands your agent might need to execute. " + + "Examples include:\n" + + "• File operations: ls, cat, grep, find, cp, mv\n" + + "• System monitoring: ps, top, df, netstat, ss\n" + + "• Process management: systemctl, service, kill\n" + + "• Network operations: ping, wget, curl\n" + + "• Data processing: awk, sed, sort, uniq\n\n" + + "What commands should your agent be able to run?"; + } + + if (lowerMessage.contains("activity") || lowerMessage.contains("activities")) { + return "Activities define what your agent can do at a high level. Examples include:\n" + + "• file_operations - reading, writing, managing files\n" + + "• system_monitoring - checking system health and metrics\n" + + "• data_processing - transforming, analyzing data\n" + + "• network_access - making network requests\n" + + "• user_management - managing user accounts\n" + + "• service_management - starting, stopping services\n\n" + + "What activities should your agent perform?"; + } + + if (lowerMessage.contains("existing") || lowerMessage.contains("policies")) { + List existingPolicies = atplPolicyService.getAllPolicies(); + if (existingPolicies.isEmpty()) { + return "No existing ATPL policies found. Would you like to create your first policy?"; + } + + StringBuilder response = new StringBuilder(); + response.append("Here are the existing ATPL policies:\n\n"); + + existingPolicies.forEach(policy -> { + response.append("**").append(policy.getPolicyId()).append("**\n"); + response.append("Version: ").append(policy.getVersion()).append("\n"); + if (policy.getDescription() != null) { + response.append("Description: ").append(policy.getDescription()).append("\n"); + } + + if (policy.getCapabilities() != null && policy.getCapabilities().getPrimitives() != null) { + response.append("Capabilities:\n"); + policy.getCapabilities().getPrimitives().forEach(cap -> { + response.append("• ").append(cap.getId()).append(": ").append(cap.getDescription()).append("\n"); + }); + } + response.append("\n"); + }); + + response.append("Would you like to create a new policy or modify an existing one?"); + return response.toString(); + } + + if (lowerMessage.contains("help") || lowerMessage.contains("what")) { + return "I'm here to help you create ATPL (Agent Trust Policy Language) policies. " + + "I can assist with:\n\n" + + "• **Policy Structure** - Understanding the basic components\n" + + "• **Endpoints** - Defining what APIs your agent can access\n" + + "• **Commands** - Specifying what system commands are allowed\n" + + "• **Activities** - Defining high-level capabilities\n" + + "• **Security** - Ensuring appropriate access controls\n\n" + + "To get started, tell me about your agent's purpose or ask about any specific aspect!"; + } + + if (lowerMessage.contains("security") || lowerMessage.contains("access")) { + return "Security is crucial for ATPL policies. Consider these aspects:\n\n" + + "• **Principle of Least Privilege** - Only grant necessary permissions\n" + + "• **Risk Assessment** - Tag high-risk capabilities appropriately\n" + + "• **Endpoint Validation** - Ensure endpoints are legitimate and necessary\n" + + "• **Command Restrictions** - Avoid dangerous commands like rm -rf, sudo\n" + + "• **Activity Boundaries** - Clearly define what activities are allowed\n\n" + + "What security considerations do you have for your agent?"; + } + + // Default response + return "I understand you want to configure an ATPL policy. I can help you define:\n" + + "• **Endpoints** - API endpoints your agent can access\n" + + "• **Commands** - System commands your agent can execute\n" + + "• **Activities** - High-level capabilities your agent needs\n\n" + + "What would you like to configure first? Or tell me more about what your agent needs to do."; + } + + public ObjectNode generateATPLPolicy(Map configuration) { + try { + ObjectNode policyNode = objectMapper.createObjectNode(); + + // Basic information + policyNode.put("version", "v0"); + policyNode.put("policy_id", (String) configuration.getOrDefault("policy_id", "generated_policy_" + System.currentTimeMillis())); + policyNode.put("description", (String) configuration.getOrDefault("description", "Generated ATPL policy from chat session")); + + // Capabilities + ObjectNode capabilitiesNode = objectMapper.createObjectNode(); + ArrayNode primitivesArray = objectMapper.createArrayNode(); + + // Extract capabilities from session context + @SuppressWarnings("unchecked") + List> capabilities = (List>) configuration.get("capabilities"); + + if (capabilities != null) { + for (Map cap : capabilities) { + ObjectNode capNode = objectMapper.createObjectNode(); + capNode.put("id", (String) cap.get("id")); + capNode.put("description", (String) cap.get("description")); + + // Add endpoints + if (cap.containsKey("endpoints")) { + ArrayNode endpointsArray = objectMapper.createArrayNode(); + @SuppressWarnings("unchecked") + List endpoints = (List) cap.get("endpoints"); + endpoints.forEach(endpointsArray::add); + capNode.set("endpoints", endpointsArray); + } + + // Add commands + if (cap.containsKey("commands")) { + ArrayNode commandsArray = objectMapper.createArrayNode(); + @SuppressWarnings("unchecked") + List commands = (List) cap.get("commands"); + commands.forEach(commandsArray::add); + capNode.set("commands", commandsArray); + } + + // Add activities + if (cap.containsKey("activities")) { + ArrayNode activitiesArray = objectMapper.createArrayNode(); + @SuppressWarnings("unchecked") + List activities = (List) cap.get("activities"); + activities.forEach(activitiesArray::add); + capNode.set("activities", activitiesArray); + } + + primitivesArray.add(capNode); + } + } else { + // Create default capability based on conversation + ObjectNode defaultCap = objectMapper.createObjectNode(); + defaultCap.put("id", "basic_access"); + defaultCap.put("description", "Basic agent access capability"); + + ArrayNode endpoints = objectMapper.createArrayNode(); + endpoints.add("/api/v1/status"); + defaultCap.set("endpoints", endpoints); + + ArrayNode commands = objectMapper.createArrayNode(); + commands.add("ls"); + commands.add("ps"); + defaultCap.set("commands", commands); + + ArrayNode activities = objectMapper.createArrayNode(); + activities.add("monitoring"); + defaultCap.set("activities", activities); + + primitivesArray.add(defaultCap); + } + + capabilitiesNode.set("primitives", primitivesArray); + policyNode.set("capabilities", capabilitiesNode); + + return policyNode; + + } catch (Exception e) { + log.error("Error generating ATPL policy: {}", e.getMessage(), e); + return objectMapper.createObjectNode(); + } + } + + public List suggestCapabilities(String userDescription) { + List suggestions = new ArrayList<>(); + + String description = userDescription.toLowerCase(); + + if (description.contains("read") || description.contains("view") || description.contains("get")) { + suggestions.add("read_access"); + } + if (description.contains("write") || description.contains("modify") || description.contains("update")) { + suggestions.add("write_access"); + } + if (description.contains("execute") || description.contains("run") || description.contains("command")) { + suggestions.add("execute_access"); + } + if (description.contains("admin") || description.contains("manage") || description.contains("control")) { + suggestions.add("admin_access"); + } + if (description.contains("network") || description.contains("connection") || description.contains("socket")) { + suggestions.add("network_access"); + } + if (description.contains("file") || description.contains("disk") || description.contains("storage")) { + suggestions.add("filesystem_access"); + } + if (description.contains("monitor") || description.contains("watch") || description.contains("observe")) { + suggestions.add("monitoring_access"); + } + if (description.contains("database") || description.contains("db") || description.contains("sql")) { + suggestions.add("database_access"); + } + + return suggestions; + } +} \ No newline at end of file diff --git a/api/src/main/resources/templates/sso/atpl/chat.html b/api/src/main/resources/templates/sso/atpl/chat.html new file mode 100644 index 00000000..c0938f68 --- /dev/null +++ b/api/src/main/resources/templates/sso/atpl/chat.html @@ -0,0 +1,555 @@ + + + + + [[${systemOptions.systemLogoName}]] - ATPL Chat Configuration + + + +
+ +
+
+
+ +
+
+

ATPL Configuration Assistant

+

Let me help you create a comprehensive Agent Trust Policy Language configuration

+
+ +
+
+ Welcome! I'm here to help you create an ATPL policy. + What kind of agent are you looking to configure? +
+
+ +
+
+ + + + +
+ +
+ + +
+
+ +
+
Assistant is thinking... +
+
+ +
+ + + +
+
+
+ + + + \ No newline at end of file diff --git a/api/src/main/resources/templates/sso/atpl/list.html b/api/src/main/resources/templates/sso/atpl/list.html index 00c1ebe8..3b7fee2c 100644 --- a/api/src/main/resources/templates/sso/atpl/list.html +++ b/api/src/main/resources/templates/sso/atpl/list.html @@ -72,7 +72,14 @@

ATPL Policy List

- + Create New Policy +
+ + Chat Assistant + + + Create New Policy + +
diff --git a/core/src/main/java/io/sentrius/sso/core/trust/Capability.java b/core/src/main/java/io/sentrius/sso/core/trust/Capability.java index fb624a1b..9d68711b 100644 --- a/core/src/main/java/io/sentrius/sso/core/trust/Capability.java +++ b/core/src/main/java/io/sentrius/sso/core/trust/Capability.java @@ -25,6 +25,12 @@ public class Capability { private List endpoints = new ArrayList<>(); @Builder.Default private List tags = new ArrayList<>(); + @Builder.Default + private List commands = new ArrayList<>(); + @Builder.Default + private List activities = new ArrayList<>(); + @Builder.Default + private List subcommands = new ArrayList<>(); public String getId() { return id; From e878eee0d38fc46cadc15b3032ee58e3f200f785 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 8 Jul 2025 00:53:30 +0000 Subject: [PATCH 3/9] Add comprehensive tests for ATPL chat service and finalize implementation Co-authored-by: phrocker <1781585+phrocker@users.noreply.github.com> --- .../sso/services/ATPLChatServiceTest.java | 160 ++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 api/src/test/java/io/sentrius/sso/services/ATPLChatServiceTest.java diff --git a/api/src/test/java/io/sentrius/sso/services/ATPLChatServiceTest.java b/api/src/test/java/io/sentrius/sso/services/ATPLChatServiceTest.java new file mode 100644 index 00000000..8a90455e --- /dev/null +++ b/api/src/test/java/io/sentrius/sso/services/ATPLChatServiceTest.java @@ -0,0 +1,160 @@ +package io.sentrius.sso.services; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; + +import io.sentrius.sso.core.services.ATPLPolicyService; +import io.sentrius.sso.core.trust.ATPLPolicy; + +@ExtendWith(MockitoExtension.class) +class ATPLChatServiceTest { + + @Mock + private ATPLPolicyService atplPolicyService; + + private ObjectMapper objectMapper; + private ATPLChatService atplChatService; + + @BeforeEach + void setUp() { + objectMapper = new ObjectMapper(); + atplChatService = new ATPLChatService(atplPolicyService, objectMapper); + } + + @Test + void testProcessATPLChatMessage_StartMessage() { + // Given + String userMessage = "I want to start creating a new ATPL policy"; + Map context = new HashMap<>(); + + // When + String response = atplChatService.processATPLChatMessage(userMessage, context); + + // Then + assertNotNull(response); + assertTrue(response.contains("I'll help you create a new ATPL policy")); + assertTrue(response.contains("File operations")); + assertTrue(response.contains("System monitoring")); + } + + @Test + void testProcessATPLChatMessage_EndpointMessage() { + // Given + String userMessage = "What endpoints should I define?"; + Map context = new HashMap<>(); + + // Mock existing policies + List existingPolicies = new ArrayList<>(); + when(atplPolicyService.getAllPolicies()).thenReturn(existingPolicies); + + // When + String response = atplChatService.processATPLChatMessage(userMessage, context); + + // Then + assertNotNull(response); + assertTrue(response.contains("endpoints")); + assertTrue(response.contains("/api/v1/")); + } + + @Test + void testProcessATPLChatMessage_CommandMessage() { + // Given + String userMessage = "What commands should I allow?"; + Map context = new HashMap<>(); + + // When + String response = atplChatService.processATPLChatMessage(userMessage, context); + + // Then + assertNotNull(response); + assertTrue(response.contains("commands")); + assertTrue(response.contains("ls")); + assertTrue(response.contains("ps")); + } + + @Test + void testSuggestCapabilities() { + // Given + String description = "I need an agent that can read files and monitor the system"; + + // When + List suggestions = atplChatService.suggestCapabilities(description); + + // Then + assertNotNull(suggestions); + assertTrue(suggestions.contains("read_access")); + assertTrue(suggestions.contains("monitoring_access")); + assertTrue(suggestions.contains("filesystem_access")); + } + + @Test + void testGenerateATPLPolicy() { + // Given + Map configuration = new HashMap<>(); + configuration.put("policy_id", "test_policy"); + configuration.put("description", "Test policy description"); + + // When + ObjectNode policy = atplChatService.generateATPLPolicy(configuration); + + // Then + assertNotNull(policy); + assertEquals("v0", policy.get("version").asText()); + assertEquals("test_policy", policy.get("policy_id").asText()); + assertEquals("Test policy description", policy.get("description").asText()); + assertTrue(policy.has("capabilities")); + assertTrue(policy.get("capabilities").has("primitives")); + } + + @Test + void testGenerateATPLPolicy_WithCapabilities() { + // Given + Map configuration = new HashMap<>(); + configuration.put("policy_id", "test_policy_with_caps"); + + List> capabilities = new ArrayList<>(); + Map capability = new HashMap<>(); + capability.put("id", "test_capability"); + capability.put("description", "Test capability"); + + List endpoints = List.of("/api/v1/test"); + List commands = List.of("ls", "ps"); + List activities = List.of("monitoring"); + + capability.put("endpoints", endpoints); + capability.put("commands", commands); + capability.put("activities", activities); + + capabilities.add(capability); + configuration.put("capabilities", capabilities); + + // When + ObjectNode policy = atplChatService.generateATPLPolicy(configuration); + + // Then + assertNotNull(policy); + assertTrue(policy.has("capabilities")); + assertTrue(policy.get("capabilities").has("primitives")); + + ObjectNode firstCapability = (ObjectNode) policy.get("capabilities").get("primitives").get(0); + assertEquals("test_capability", firstCapability.get("id").asText()); + assertEquals("Test capability", firstCapability.get("description").asText()); + assertTrue(firstCapability.has("endpoints")); + assertTrue(firstCapability.has("commands")); + assertTrue(firstCapability.has("activities")); + } +} \ No newline at end of file From 5f2616067c33a2bff4bf71aa831b967d2b8078e3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 8 Jul 2025 01:32:32 +0000 Subject: [PATCH 4/9] Fix template reference error in ATPL chat interface Co-authored-by: phrocker <1781585+phrocker@users.noreply.github.com> --- api/src/main/resources/templates/sso/atpl/chat.html | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/api/src/main/resources/templates/sso/atpl/chat.html b/api/src/main/resources/templates/sso/atpl/chat.html index c0938f68..e7842d51 100644 --- a/api/src/main/resources/templates/sso/atpl/chat.html +++ b/api/src/main/resources/templates/sso/atpl/chat.html @@ -253,10 +253,11 @@ -
- -
-
+
+
+
+
+
@@ -316,6 +317,7 @@

ATPL Configuration Assistant

+
+
+
@@ -266,14 +306,14 @@

ATPL Configuration Assistant

Let me help you create a comprehensive Agent Trust Policy Language configuration

-
+
Welcome! I'm here to help you create an ATPL policy. What kind of agent are you looking to configure?
-
+
- -
@@ -319,66 +359,105 @@

ATPL Configuration Assistant

- \ No newline at end of file diff --git a/core/src/main/java/io/sentrius/sso/core/dto/capabilities/AccessLimitations.java b/core/src/main/java/io/sentrius/sso/core/dto/capabilities/AccessLimitations.java index b65789f5..02452eb0 100644 --- a/core/src/main/java/io/sentrius/sso/core/dto/capabilities/AccessLimitations.java +++ b/core/src/main/java/io/sentrius/sso/core/dto/capabilities/AccessLimitations.java @@ -8,9 +8,11 @@ import io.sentrius.sso.core.model.security.enums.SystemOperationsEnum; import io.sentrius.sso.core.model.security.enums.UserAccessEnum; import io.sentrius.sso.core.model.security.enums.ZeroTrustAccessTokenEnum; +import lombok.AllArgsConstructor; import lombok.Builder; import lombok.Data; import lombok.Getter; +import lombok.NoArgsConstructor; import lombok.Setter; /** @@ -20,6 +22,8 @@ @Data @Getter @Setter +@NoArgsConstructor +@AllArgsConstructor public class AccessLimitations { private String notificationMessage; private IdentityType[] allowedIdentityTypes; diff --git a/core/src/main/java/io/sentrius/sso/core/dto/capabilities/EndpointDescriptor.java b/core/src/main/java/io/sentrius/sso/core/dto/capabilities/EndpointDescriptor.java index c9813088..cf6fec30 100644 --- a/core/src/main/java/io/sentrius/sso/core/dto/capabilities/EndpointDescriptor.java +++ b/core/src/main/java/io/sentrius/sso/core/dto/capabilities/EndpointDescriptor.java @@ -2,10 +2,13 @@ import java.util.List; import java.util.Map; +import lombok.AllArgsConstructor; import lombok.Builder; import lombok.Data; import lombok.Getter; +import lombok.NoArgsConstructor; import lombok.Setter; +import lombok.ToString; /** * Represents a unified descriptor for both REST API endpoints and Verb methods. @@ -15,6 +18,9 @@ @Data @Getter @Setter +@ToString +@NoArgsConstructor +@AllArgsConstructor public class EndpointDescriptor { private String name; private String description; diff --git a/core/src/main/java/io/sentrius/sso/core/dto/capabilities/ParameterDescriptor.java b/core/src/main/java/io/sentrius/sso/core/dto/capabilities/ParameterDescriptor.java index a108a1d1..c90df220 100644 --- a/core/src/main/java/io/sentrius/sso/core/dto/capabilities/ParameterDescriptor.java +++ b/core/src/main/java/io/sentrius/sso/core/dto/capabilities/ParameterDescriptor.java @@ -1,8 +1,10 @@ package io.sentrius.sso.core.dto.capabilities; +import lombok.AllArgsConstructor; import lombok.Builder; import lombok.Data; import lombok.Getter; +import lombok.NoArgsConstructor; import lombok.Setter; /** @@ -12,10 +14,12 @@ @Data @Getter @Setter +@NoArgsConstructor +@AllArgsConstructor public class ParameterDescriptor { private String name; private String description; - private Class type; + private String type; private boolean required; private Object defaultValue; private String source; // "PATH", "QUERY", "BODY", "HEADER", "METHOD_PARAM" diff --git a/core/src/main/java/io/sentrius/sso/core/dto/ztat/AgentExecution.java b/core/src/main/java/io/sentrius/sso/core/dto/ztat/AgentExecution.java index 5e01c3df..81789c1b 100644 --- a/core/src/main/java/io/sentrius/sso/core/dto/ztat/AgentExecution.java +++ b/core/src/main/java/io/sentrius/sso/core/dto/ztat/AgentExecution.java @@ -30,4 +30,9 @@ public void addMessages(List messages) { this.messages.addAll(messages); } + + public void addMessages(Message message) { + this.messages.add(message); + + } } diff --git a/core/src/main/java/io/sentrius/sso/core/model/verbs/ListInterpreter.java b/core/src/main/java/io/sentrius/sso/core/model/verbs/ListInterpreter.java index 050784ac..61ff0263 100644 --- a/core/src/main/java/io/sentrius/sso/core/model/verbs/ListInterpreter.java +++ b/core/src/main/java/io/sentrius/sso/core/model/verbs/ListInterpreter.java @@ -20,7 +20,7 @@ public List interpret(Map input) throws Exception { TypeReference> typeRef = new TypeReference<>() {}; return JsonUtil.convertArrayNodeToList(node,typeRef); } else { - throw new IllegalArgumentException("Invalid input type " + input.get("verb.response.type")); + return null; } } @@ -30,13 +30,13 @@ private List interpretList(Map input) { var field = input.get("verb.response.map.key"); if (field == null) { - throw new IllegalArgumentException("Input response does not contain required fields"); + return null; } var object = input.get(field); if (object instanceof List){ return (List) object; } - throw new IllegalArgumentException("Input response does not contain required fields"); + return null; } } diff --git a/core/src/main/java/io/sentrius/sso/core/services/agents/AgentClientService.java b/core/src/main/java/io/sentrius/sso/core/services/agents/AgentClientService.java index 712dc624..8b1b55ea 100644 --- a/core/src/main/java/io/sentrius/sso/core/services/agents/AgentClientService.java +++ b/core/src/main/java/io/sentrius/sso/core/services/agents/AgentClientService.java @@ -11,6 +11,7 @@ import io.sentrius.sso.core.dto.AgentCommunicationDTO; import io.sentrius.sso.core.dto.AgentHeartbeatDTO; import io.sentrius.sso.core.dto.AgentRegistrationDTO; +import io.sentrius.sso.core.dto.capabilities.EndpointDescriptor; import io.sentrius.sso.core.dto.ztat.AgentExecution; import io.sentrius.sso.core.dto.ztat.AtatRequest; import io.sentrius.sso.core.dto.ztat.TokenDTO; @@ -207,4 +208,32 @@ public String getCallbackUrl() { return callbackUrl; } + public List getAvailableEndpoints(TokenDTO token) throws ZtatException, JsonProcessingException { + String url = "/api/v1/capabilities/endpoints"; + Object response = zeroTrustClientService.callGetOnApi(token, url); + if (response instanceof String str) { + return JsonUtil.MAPPER.readValue(str, new TypeReference>() {}); + } else if (response instanceof List list) { + // Already deserialized + return list.stream() + .map(item -> JsonUtil.MAPPER.convertValue(item, EndpointDescriptor.class)) + .toList(); + } + return List.of(); + } + + public List getAvailableVerbs(TokenDTO token) throws ZtatException, JsonProcessingException { + String url = "/api/v1/capabilities/verbs"; + Object response = zeroTrustClientService.callGetOnApi(token, url); + if (response instanceof String str) { + return JsonUtil.MAPPER.readValue(str, new TypeReference>() {}); + } else if (response instanceof List list) { + // Already deserialized + return list.stream() + .map(item -> JsonUtil.MAPPER.convertValue(item, EndpointDescriptor.class)) + .toList(); + } + return List.of(); + } + } diff --git a/core/src/main/java/io/sentrius/sso/core/services/capabilities/EndpointScanningService.java b/core/src/main/java/io/sentrius/sso/core/services/capabilities/EndpointScanningService.java index 2ece69d2..00394ac6 100644 --- a/core/src/main/java/io/sentrius/sso/core/services/capabilities/EndpointScanningService.java +++ b/core/src/main/java/io/sentrius/sso/core/services/capabilities/EndpointScanningService.java @@ -284,7 +284,7 @@ private List extractRestParameters(Method method) { parameters.add(ParameterDescriptor.builder() .name(name) - .type(parameter.getType()) + .type(parameter.getType().getCanonicalName()) .required(required) .source(source) .build()); @@ -307,7 +307,7 @@ private List extractVerbParameters(Method method, Verb verb parameters.add(ParameterDescriptor.builder() .name(parameter.getName()) .description(description) - .type(parameter.getType()) + .type(parameter.getType().getCanonicalName()) .required(true) // Assume required for verb parameters .source("METHOD_PARAM") .build()); diff --git a/core/src/main/java/io/sentrius/sso/core/services/security/JwtUtil.java b/core/src/main/java/io/sentrius/sso/core/services/security/JwtUtil.java index dde006eb..b62b10f2 100644 --- a/core/src/main/java/io/sentrius/sso/core/services/security/JwtUtil.java +++ b/core/src/main/java/io/sentrius/sso/core/services/security/JwtUtil.java @@ -95,10 +95,7 @@ public static String extractKid(String jwt) { jwt = jwt.substring(7); } String[] parts = jwt.split("\\."); - log.info("JWT parts: header={}, payload={}, signaturePresent={}", - parts.length > 0 ? parts[0] : "null", - parts.length > 1 ? parts[1] : "null", - parts.length == 3); + if (parts.length != 3) { throw new IllegalArgumentException("Invalid JWT token format"); } diff --git a/dataplane/src/main/java/io/sentrius/sso/core/model/security/AccessControlAspect.java b/dataplane/src/main/java/io/sentrius/sso/core/model/security/AccessControlAspect.java index 888d4cd9..a9a4dd6f 100644 --- a/dataplane/src/main/java/io/sentrius/sso/core/model/security/AccessControlAspect.java +++ b/dataplane/src/main/java/io/sentrius/sso/core/model/security/AccessControlAspect.java @@ -70,16 +70,20 @@ public class AccessControlAspect { allowedEndpoints.add("/api/v1/zerotrust/accesstoken/status"); allowedEndpoints.add("/api/v1/zerotrust/accesstoken/jwt/verify"); allowedEndpoints.add("/api/v1/agent/bootstrap/register"); + allowedEndpoints.add("/api/v1/capabilities/endpoints"); + allowedEndpoints.add("/api/v1/capabilities/verbs"); } Tracer tracer = GlobalOpenTelemetry.getTracer("io.sentrius.sso"); private boolean isAllowedEndpoint(String endpoint) { for (String allowedEndpoint : allowedEndpoints) { + log.info("Checking if endpoint {} matches {}", endpoint, allowedEndpoint); if (endpoint.startsWith(allowedEndpoint)) { return true; } } + log.info("Endpoint {} doesn't match", endpoint); return false; } diff --git a/dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java b/dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java index ca397061..baf59342 100644 --- a/dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java +++ b/dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java @@ -44,6 +44,7 @@ import org.springframework.scheduling.annotation.Async; import org.springframework.scheduling.annotation.Scheduled; import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; import org.springframework.web.client.RestTemplate; @Service @@ -84,6 +85,7 @@ public AgentService( this.provenanceKafkaProducer = provenanceKafkaProducer; } + @Transactional public void recordHeartbeat(String agentId, String name, AgentHeartbeatDTO heartbeatDTO) { AgentHeartbeat heartbeat = repository.findByAgentId(agentId) .orElse(new AgentHeartbeat()); diff --git a/integration-proxy/src/main/java/io/sentrius/sso/controllers/api/OpenAIProxyController.java b/integration-proxy/src/main/java/io/sentrius/sso/controllers/api/OpenAIProxyController.java index 4ee6dd07..0a83ee29 100644 --- a/integration-proxy/src/main/java/io/sentrius/sso/controllers/api/OpenAIProxyController.java +++ b/integration-proxy/src/main/java/io/sentrius/sso/controllers/api/OpenAIProxyController.java @@ -180,20 +180,35 @@ public ResponseEntity chat(@RequestHeader("Authorization") String token, Span span = tracer.spanBuilder("AgentToAgentCommunication").startSpan(); + int retries = 2; try (Scope scope = span.makeCurrent()) { - var resp = endpoint.sample(RawConversationRequest.builder().request(chatRequest).build()); - span.setAttribute("communication.id", comm.get().getId().toString()); - span.setAttribute("source.agent", operatingUser.getUsername()); - span.setAttribute("target.agent", "SYSTEM"); - span.setAttribute("message.type", "interpretation_request"); - return ResponseEntity.ok(resp); + HttpException httpException = null; + do { + try { + var resp = endpoint.sample(RawConversationRequest.builder().request(chatRequest).build()); + span.setAttribute("communication.id", comm.get().getId().toString()); + span.setAttribute("source.agent", operatingUser.getUsername()); + span.setAttribute("target.agent", "SYSTEM"); + span.setAttribute("message.type", "interpretation_request"); + return ResponseEntity.ok(resp); + }catch(HttpException e){ + if (e.getMessage().contains("timeout")) { + httpException = e; + } else { + throw e; + } + } + } while(retries-- > 0); + if (null != httpException) { + throw httpException; + } } catch (ExecutionException | InterruptedException e) { throw new RuntimeException(e); } finally { span.end(); } - + return null; } @PostMapping("/justify") diff --git a/llm-dataplane/src/main/java/io/sentrius/sso/genai/GenerativeAPI.java b/llm-dataplane/src/main/java/io/sentrius/sso/genai/GenerativeAPI.java index d1c01fc9..5f767633 100644 --- a/llm-dataplane/src/main/java/io/sentrius/sso/genai/GenerativeAPI.java +++ b/llm-dataplane/src/main/java/io/sentrius/sso/genai/GenerativeAPI.java @@ -1,6 +1,7 @@ package io.sentrius.sso.genai; import java.io.IOException; +import java.time.Duration; import java.util.Objects; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; @@ -42,12 +43,17 @@ public GenerativeAPI(TokenProvider authToken, OkHttpClient client) { } public GenerativeAPI(TokenProvider authToken) { - this(authToken, new OkHttpClient()); + + this(authToken, new OkHttpClient.Builder() + .connectTimeout(Duration.ofSeconds(15)) + .readTimeout(Duration.ofSeconds(30)) + .writeTimeout(Duration.ofSeconds(15)) + .build()); } /** * Builds the request body for an API endpoint request. - * + *a * @param request the API endpoint request * @return the request body as a string */ @@ -84,6 +90,7 @@ public String sample(final ApiEndPointRequest apiRequest) throws HttpException { throw new HttpException(response.code(), response.body().string()); } } else { + log.info("body is {}" , response.body()); return response.body().string(); } } catch (IOException e) { diff --git a/sentrius-chart-launcher/templates/configmap.yaml b/sentrius-chart-launcher/templates/configmap.yaml index 169431fd..cd2493a8 100644 --- a/sentrius-chart-launcher/templates/configmap.yaml +++ b/sentrius-chart-launcher/templates/configmap.yaml @@ -16,6 +16,211 @@ data: "terminalSummaryForLLM": "", "responseForUser": "" } + chat-atpl-helper.yaml: | + description: "Agent that handles logs and OpenAI access." + context: | + You are a helpful agent that helps define a a trust policy. Guide the user through the process of defining a trust policy for an agent. + Here is the schema for the Agent Trust Policy Language (ATPL). If no version or policy_id given generate them. + Most fields are optional, but the minimum will be capabilities. primitives are simple endpoint connections. + Focus on getting those filled in. Following this list will be a message containing all endpoints of the system. + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://sentrius.io/schemas/atpl.schema.json", + "title": "Agent Trust Policy Language (ATPL)", + "type": "object", + "required": ["version", "policy_id", "capabilities"], + "properties": { + "version": { + "type": "string", + "pattern": "^v[0-9]+$" + }, + "policy_id": { + "type": "string" + }, + "description": { + "type": "string" + }, + "match": { + "type": "object", + "properties": { + "agent_tags": { + "type": "array", + "items": { "type": "string" } + } + } + }, + "identity": { + "type": "object", + "properties": { + "issuer": { "type": "string" }, + "subject_prefix": { "type": "string" }, + "mfa_required": { "type": "boolean" }, + "certificate_authority": { "type": "string" } + } + }, + "provenance": { + "type": "object", + "properties": { + "source": { "type": "string" }, + "signature_required": { "type": "boolean" }, + "approved_committers": { + "type": "array", + "items": { "type": "string" } + } + } + }, + "runtime": { + "type": "object", + "properties": { + "enclave_required": { "type": "boolean" }, + "attestation_type": { "type": "string" }, + "verified_at_boot": { "type": "boolean" } + } + }, + "behavior": { + "type": "object", + "properties": { + "minimum_positive_runs": { "type": "integer" }, + "max_incidents": { "type": "integer" }, + "incident_types": { + "type": "object", + "properties": { + "denylist": { + "type": "array", + "items": { "type": "string" } + } + } + } + } + }, + "trust_score": { + "type": "object", + "properties": { + "minimum": { "type": "integer" }, + "weightings": { + "type": "object", + "properties": { + "identity": { "type": "number" }, + "provenance": { "type": "number" }, + "runtime": { "type": "number" }, + "behavior": { "type": "number" } + } + } + } + }, + "capabilities": { + "type": "object", + "required": ["primitives", "composed"], + "properties": { + "primitives": { + "type": "array", + "items": { + "type": "object", + "required": ["id", "description"], + "properties": { + "id": { "type": "string" }, + "description": { "type": "string" }, + "endpoints": { + "type": "array", + "items": { "type": "string" } + }, + "commands": { + "type": "array", + "items": { "type": "string" } + }, + "subcommands": { + "type": "array", + "items": { "type": "string" } + }, + "activities": { + "type": "array", + "items": { "type": "string" } + }, + "tags": { + "type": "array", + "items": { "type": "string" } + } + } + } + }, + "composed": { + "type": "array", + "items": { + "type": "object", + "required": ["id", "includes"], + "properties": { + "id": { "type": "string" }, + "includes": { + "type": "array", + "items": { "type": "string" } + }, + "tags": { + "type": "array", + "items": { "type": "string" } + }, + "endpoints": { + "type": "array", + "items": { "type": "string" } + }, + "commands": { + "type": "array", + "items": { "type": "string" } + }, + "subcommands": { + "type": "array", + "items": { "type": "string" } + }, + "activities": { + "type": "array", + "items": { "type": "string" } + } + } + } + } + } + } + , + "actions": { + "type": "object", + "properties": { + "on_failure": { + "type": "string", + "enum": ["deny", "log", "alert"] + }, + "on_success": { + "type": "string", + "enum": ["allow", "log"] + }, + "on_marginal": { + "oneOf": [ + {"type": "string", "enum": ["require_ztat", "log"]}, + { + "type": "object", + "properties": { + "action": {"type": "string", "enum": ["require_ztat"]}, + "ztat_provider": {"type": "string"} + }, + "required": ["action"] + } + ] + } + } + }, + "ztat": { + "type": "object", + "properties": { + "provider": {"type": "string"}, + "ttl": {"type": "string"}, + "approved_issuers": { + "type": "array", + "items": {"type": "string"} + }, + "key_binding": {"type": "string"}, + "approval_required": {"type": "boolean"} + } + } + } + } launcher.properties: | spring.main.web-application-type=servlet spring.thymeleaf.enabled=false From 844895202e7b94e70173fe05df2694b3bd557d11 Mon Sep 17 00:00:00 2001 From: Marc Parisi Date: Wed, 9 Jul 2025 18:55:25 -0400 Subject: [PATCH 6/9] intermediate commit --- .local.env | 8 +- .local.env.bak | 8 +- .../launcher/api/AgentLauncherController.java | 10 +- .../launcher/service/PodLauncherService.java | 44 +- .../config/AgentWebSocketProxyHandler.java | 24 +- .../sentrius/sso/config/SecurityConfig.java | 4 - .../analysis/agents/agents/AgentVerb.java | 2 + .../analysis/agents/agents/PromptBuilder.java | 39 +- .../analysis/agents/agents/VerbRegistry.java | 1 + .../interpreters/AgentContextInterpreter.java | 22 + .../interpreters/ObjectNodeInterpreter.java | 37 + .../interpreters/StringInterpreter.java | 22 + .../interpreters/StringToAtplInterpreter.java | 23 + .../analysis/agents/verbs/AgentVerbs.java | 131 +++- .../analysis/agents/verbs/AtplVerbs.java | 57 +- .../analysis/agents/verbs/ChatVerbs.java | 173 +++-- .../analysis/agents/verbs/ExampleFactory.java | 41 ++ .../analysis/api/websocket/ChatWSHandler.java | 9 +- .../analysis/model/TerminalResponse.java | 3 + .../sentrius/agent/config/SecurityConfig.java | 2 - ai-agent/src/main/resources/atpl-schema.json | 199 +++++ ai-agent/src/main/resources/chat-helper.json | 4 + .../src/main/resources/terminal-helper.json | 6 +- .../controllers/api/AgentApiController.java | 74 +- .../api/AgentBootstrapController.java | 38 +- .../controllers/api/HostApiController.java | 2 +- .../controllers/api/SystemApiController.java | 16 +- .../sso/controllers/view/AgentController.java | 7 + .../db/migration/V18_agent_launch.sql | 17 + .../main/resources/static/js/add_system.js | 14 +- .../templates/sso/agents/design_chat.html | 692 ++++++++++++++++++ .../resources/templates/sso/dashboard.html | 329 ++++----- .../templates/sso/enclaves/list_servers.html | 2 +- .../sso/core/dto/AgentRegistrationDTO.java | 2 + .../sso/core/dto/agents/AgentContextDTO.java | 29 + .../dto/agents/AgentContextRequestDTO.java | 23 + .../dto/agents/AgentLaunchResponseDTO.java | 11 + .../sentrius/sso/core/model/verbs/Verb.java | 4 + .../services/agents/AgentClientService.java | 43 ++ .../agents/ZeroTrustClientService.java | 7 +- .../sentrius/sso/core/trust/ATPLPolicy.java | 2 + .../sso/core/config/SystemOptions.java | 3 +- .../sso/core/model/agents/AgentContext.java | 46 ++ .../sso/core/model/agents/AgentLaunch.java | 38 + .../repository/AgentContextRepository.java | 9 + .../repository/AgentLaunchRepository.java | 8 + .../services/agents/AgentContextService.java | 31 + .../services/agents/AgentLaunchService.java | 32 + .../templates/configmap.yaml | 220 +----- 49 files changed, 2077 insertions(+), 491 deletions(-) create mode 100644 ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/AgentContextInterpreter.java create mode 100644 ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/ObjectNodeInterpreter.java create mode 100644 ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/StringInterpreter.java create mode 100644 ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/StringToAtplInterpreter.java create mode 100644 ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ExampleFactory.java create mode 100644 ai-agent/src/main/resources/atpl-schema.json create mode 100644 api/src/main/resources/db/migration/V18_agent_launch.sql create mode 100644 api/src/main/resources/templates/sso/agents/design_chat.html create mode 100644 core/src/main/java/io/sentrius/sso/core/dto/agents/AgentContextDTO.java create mode 100644 core/src/main/java/io/sentrius/sso/core/dto/agents/AgentContextRequestDTO.java create mode 100644 core/src/main/java/io/sentrius/sso/core/dto/agents/AgentLaunchResponseDTO.java create mode 100644 dataplane/src/main/java/io/sentrius/sso/core/model/agents/AgentContext.java create mode 100644 dataplane/src/main/java/io/sentrius/sso/core/model/agents/AgentLaunch.java create mode 100644 dataplane/src/main/java/io/sentrius/sso/core/repository/AgentContextRepository.java create mode 100644 dataplane/src/main/java/io/sentrius/sso/core/repository/AgentLaunchRepository.java create mode 100644 dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentContextService.java create mode 100644 dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentLaunchService.java diff --git a/.local.env b/.local.env index decdf6df..d6e5e1f1 100644 --- a/.local.env +++ b/.local.env @@ -1,8 +1,8 @@ -SENTRIUS_VERSION=1.1.233 +SENTRIUS_VERSION=1.1.246 SENTRIUS_SSH_VERSION=1.1.36 SENTRIUS_KEYCLOAK_VERSION=1.1.48 SENTRIUS_AGENT_VERSION=1.1.35 -SENTRIUS_AI_AGENT_VERSION=1.1.85 +SENTRIUS_AI_AGENT_VERSION=1.1.124 LLMPROXY_VERSION=1.0.49 -LAUNCHER_VERSION=1.0.55 -AGENTPROXY_VERSION=1.0.67 \ No newline at end of file +LAUNCHER_VERSION=1.0.62 +AGENTPROXY_VERSION=1.0.69 \ No newline at end of file diff --git a/.local.env.bak b/.local.env.bak index decdf6df..d6e5e1f1 100644 --- a/.local.env.bak +++ b/.local.env.bak @@ -1,8 +1,8 @@ -SENTRIUS_VERSION=1.1.233 +SENTRIUS_VERSION=1.1.246 SENTRIUS_SSH_VERSION=1.1.36 SENTRIUS_KEYCLOAK_VERSION=1.1.48 SENTRIUS_AGENT_VERSION=1.1.35 -SENTRIUS_AI_AGENT_VERSION=1.1.85 +SENTRIUS_AI_AGENT_VERSION=1.1.124 LLMPROXY_VERSION=1.0.49 -LAUNCHER_VERSION=1.0.55 -AGENTPROXY_VERSION=1.0.67 \ No newline at end of file +LAUNCHER_VERSION=1.0.62 +AGENTPROXY_VERSION=1.0.69 \ No newline at end of file diff --git a/agent-launcher/src/main/java/io/sentrius/agent/launcher/api/AgentLauncherController.java b/agent-launcher/src/main/java/io/sentrius/agent/launcher/api/AgentLauncherController.java index d44a86d1..67c73712 100644 --- a/agent-launcher/src/main/java/io/sentrius/agent/launcher/api/AgentLauncherController.java +++ b/agent-launcher/src/main/java/io/sentrius/agent/launcher/api/AgentLauncherController.java @@ -60,9 +60,17 @@ public ResponseEntity deleteAgent(@RequestParam(name="agentId") String a podLauncherService.deleteAgentById(agentId); return ResponseEntity.ok("Shutdown triggered"); } catch (Exception e) { - e.printStackTrace(); return ResponseEntity.status(500).body("Shutdown failed: " + e.getMessage()); } } + @GetMapping("/status") + public ResponseEntity getAgentStatus(@RequestParam(name="agentId") String agentId) { + try { + return ResponseEntity.ok(podLauncherService.statusById(agentId) ); + } catch (Exception e) { + return ResponseEntity.status(500).body("Status failed: " + e.getMessage()); + } + } + } diff --git a/agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java b/agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java index e08409a9..85050124 100644 --- a/agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java +++ b/agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java @@ -102,7 +102,7 @@ public void deleteAgentById(String agentId) throws Exception { } catch (Exception ex) { log.warn("Service not found or already deleted: {}", ex.getMessage()); } - }else { + } else { log.info("Not Deleting pod: {}", podName); } @@ -113,9 +113,49 @@ public void deleteAgentById(String agentId) throws Exception { } + } + + + public String statusById(String agentId) throws Exception { + // Delete all pods with this agentId label + var pods = coreV1Api.listNamespacedPod( + agentNamespace + ).execute().getItems(); + + for (V1Pod pod : pods) { + + var labels = pod.getMetadata().getLabels(); + var podName = pod.getMetadata().getName(); + + Matcher matcher = pattern.matcher(agentId); + + if (matcher.matches() && labels != null && labels.containsKey("agentId")) { + String name = matcher.group(1); + + var value = labels.get("agentId"); + if (value.equals(name)) { + // get pod status + // + V1PodStatus status = pod.getStatus(); + if (status == null) { + log.warn("Pod {} has no status information", podName); + return "Unknown"; + } + return status.getPhase(); // e.g., "Running", "Pending", "Failed", "Succeeded" + + } + + + } + + + } + return "NotFound"; + } + + - } public V1Pod launchAgentPod(AgentRegistrationDTO agent) throws Exception { diff --git a/agent-proxy/src/main/java/io/sentrius/sso/config/AgentWebSocketProxyHandler.java b/agent-proxy/src/main/java/io/sentrius/sso/config/AgentWebSocketProxyHandler.java index ecbcfa91..c801c5a7 100644 --- a/agent-proxy/src/main/java/io/sentrius/sso/config/AgentWebSocketProxyHandler.java +++ b/agent-proxy/src/main/java/io/sentrius/sso/config/AgentWebSocketProxyHandler.java @@ -123,22 +123,38 @@ public Mono handle(WebSocketSession clientSession) { }) .as(clientSession::send) .doOnSuccess(aVoid -> log.info("agent -> client completed gracefully")) // Corrected for Mono - .doOnError(e -> log.error("Error in agent -> client stream", e)) + .doOnError(e -> { + log.error("Error in agent -> client stream", e); + sessionManager.unregister(agentSession.getId()); + }) .onErrorResume(e -> { + sessionManager.unregister(agentSession.getId()); log.error("Agent to client stream error, closing agent session.", e); return agentSession.close().then(Mono.empty()); }) .doFinally(sig -> log.info("Agent to client stream finalized: {}", sig)); return Mono.when(clientToAgent, agentToClient) - .doOnTerminate(() -> log.info("WebSocket proxy connection terminated (client and agent streams completed/cancelled)")) - .doOnError(e -> log.error("Overall proxy connection failed", e)) + .doOnTerminate(() -> { + log.info("WebSocket proxy connection terminated (client and agent " + + "streams completed/cancelled)"); + sessionManager.unregister(agentSession.getId()); + + }) + .doOnError(e -> { + log.error("Overall proxy connection failed", e); + sessionManager.unregister(agentSession.getId()); + + }) .doFinally(sig -> { sessionManager.unregister(finalSessionId); log.info("WebSocket proxy stream closed completely: {}. Final session ID: {}", sig, finalSessionId); }); } - ).doOnError(e -> log.error("Failed to establish proxy connection", e)); + ).doOnError(e -> { + log.error("Failed to establish proxy connection", e); + sessionManager.unregister(finalSessionId); + }); } catch (Exception ex) { diff --git a/agent-proxy/src/main/java/io/sentrius/sso/config/SecurityConfig.java b/agent-proxy/src/main/java/io/sentrius/sso/config/SecurityConfig.java index 4595af51..17cf55df 100644 --- a/agent-proxy/src/main/java/io/sentrius/sso/config/SecurityConfig.java +++ b/agent-proxy/src/main/java/io/sentrius/sso/config/SecurityConfig.java @@ -55,10 +55,6 @@ private ReactiveJwtAuthenticationConverter grantedAuthoritiesExtractor() { converter.setJwtGrantedAuthoritiesConverter(jwt -> { Collection authorities = new JwtGrantedAuthoritiesConverter().convert(jwt); - log.info("JWT Claims: {}", jwt.getClaims()); - - String username = jwt.getClaimAsString("preferred_username"); - String email = jwt.getClaimAsString("email"); return Flux.fromIterable(authorities); }); diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/AgentVerb.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/AgentVerb.java index 523997ba..aeef6f6a 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/AgentVerb.java +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/AgentVerb.java @@ -24,4 +24,6 @@ public class AgentVerb { @Builder.Default Class outputInterpreter = DefaultInterpreter.class; Class inputInterpreter = DefaultInterpreter.class; + + private String exampleJson = ""; } \ No newline at end of file diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/PromptBuilder.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/PromptBuilder.java index 017e1856..12c7a62a 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/PromptBuilder.java +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/PromptBuilder.java @@ -3,6 +3,8 @@ import java.lang.reflect.Method; import java.util.Arrays; import java.util.stream.Collectors; +import io.sentrius.agent.analysis.agents.verbs.ExampleFactory; +import io.sentrius.sso.core.utils.JsonUtil; /** * The `PromptBuilder` class is responsible for constructing a prompt string @@ -61,21 +63,50 @@ public String buildPrompt(boolean applyInstructions) " }\n" + " ]\n" + "}\n"); - + } // Append the list of available verbs - prompt.append("Available Verbs:\n"); + prompt.append("Verb operations:\n"); // Iterate through the verbs in the registry and append their details verbRegistry.getVerbs().forEach((name, verb) -> { prompt.append("- ").append(name); prompt.append(" (").append(buildMethodSignature(verb.getMethod())).append(") - "); prompt.append(verb.getDescription()).append("\n"); + // Optionally generate example params based on arg1 class + Class[] paramTypes = verb.getMethod().getParameterTypes(); + + if (paramTypes.length > 1 && !paramTypes[1].equals(Void.class)) { + var paramName = verb.getMethod().getParameters()[1].getName(); + Object example = ExampleFactory.createExample(paramName, paramTypes[1]); // create a stub from + // your DTO + try { + if (verb.getExampleJson() != null && !verb.getExampleJson().isEmpty()) { + prompt.append(" Example arg1: ").append(verb.getExampleJson()).append("\n"); + } else if (example != null) { + // Serialize the example object to JSON + String exampleJson = JsonUtil.MAPPER.writeValueAsString(example); + prompt.append(" Example arg1: ").append(exampleJson).append("\n"); + } + + } catch (Exception e) { + prompt.append(" Example params: [unavailable due to serialization error]\n"); + } + } else { + prompt.append(" Example params: {}\n"); + } }); - } return prompt.toString(); } + public static String indent(String input, int spaces) { + String indent = " ".repeat(spaces); + return Arrays.stream(input.split("\n")) + .map(line -> indent + line) + .collect(Collectors.joining("\n")); + } + + /** * Builds a method signature string for a given method. * @@ -84,6 +115,8 @@ public String buildPrompt(boolean applyInstructions) */ private String buildMethodSignature(Method method) { return Arrays.stream(method.getParameters()) + .filter( p -> !p.getType().getSimpleName().equalsIgnoreCase("TokenDTO") && + !p.getType().getSimpleName().equalsIgnoreCase("AgentExecution")) .map(p -> p.getName() + ": " + p.getType().getSimpleName()) .collect(Collectors.joining(", ")); } diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/VerbRegistry.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/VerbRegistry.java index 54426470..c9b0e963 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/VerbRegistry.java +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/VerbRegistry.java @@ -93,6 +93,7 @@ public void scanClasspath() { .name(name) .description(annotation.description()) .method(method) + .exampleJson(annotation.exampleJson()) .requiresTokenManagement(annotation.requiresTokenManagement()) .returnType(annotation.returnType()) .outputInterpreter(annotation.outputInterpreter()) diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/AgentContextInterpreter.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/AgentContextInterpreter.java new file mode 100644 index 00000000..1743671d --- /dev/null +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/AgentContextInterpreter.java @@ -0,0 +1,22 @@ +package io.sentrius.agent.analysis.agents.interpreters; + +import java.util.Map; +import io.sentrius.sso.core.dto.agents.AgentContextDTO; +import io.sentrius.sso.core.model.verbs.InputInterpreterIfc; +import io.sentrius.sso.core.utils.JsonUtil; + + +public class AgentContextInterpreter implements InputInterpreterIfc { + @Override + public AgentContextDTO interpret(Map input) throws Exception { + Object agentContextObj = input.get("agentContext"); + Object arg1Obj = input.get("arg1"); + + if (agentContextObj != null) { + return JsonUtil.MAPPER.convertValue(agentContextObj, AgentContextDTO.class); + } else if (arg1Obj != null) { + return JsonUtil.MAPPER.convertValue(arg1Obj, AgentContextDTO.class); + } + return null; + } +} diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/ObjectNodeInterpreter.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/ObjectNodeInterpreter.java new file mode 100644 index 00000000..1f2e56ee --- /dev/null +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/ObjectNodeInterpreter.java @@ -0,0 +1,37 @@ +package io.sentrius.agent.analysis.agents.interpreters; + +import java.util.HashMap; +import java.util.Map; +import com.fasterxml.jackson.databind.node.ObjectNode; +import io.sentrius.sso.core.model.verbs.InputInterpreterIfc; +import io.sentrius.sso.core.model.verbs.OutputInterpreterIfc; +import io.sentrius.sso.core.model.verbs.VerbResponse; +import io.sentrius.sso.core.utils.JsonUtil; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +public class ObjectNodeInterpreter implements InputInterpreterIfc, OutputInterpreterIfc { + @Override + public ObjectNode interpret(Map input) throws Exception { + + Object arg1Value = input.get("arg1"); + if (arg1Value == null) { + throw new IllegalArgumentException("Missing 'arg1' field in arguments map."); + } + + // Convert the value under "arg1" to an ObjectNode + return JsonUtil.MAPPER.valueToTree(arg1Value); + } + + @Override + public Map interpret(VerbResponse input) throws Exception { + if (input.getResponse() instanceof ObjectNode) { + Map responseMap = new HashMap<>(); + ((ObjectNode)input.getResponse()).fieldNames().forEachRemaining( x ->{ + responseMap.put(x, ((ObjectNode)input.getResponse()).get(x).asText() ); + }); + return responseMap; + } + return Map.of(); + } +} diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/StringInterpreter.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/StringInterpreter.java new file mode 100644 index 00000000..996a3bc9 --- /dev/null +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/StringInterpreter.java @@ -0,0 +1,22 @@ +package io.sentrius.agent.analysis.agents.interpreters; + +import java.util.Map; +import io.sentrius.sso.core.model.verbs.InputInterpreterIfc; +import io.sentrius.sso.core.trust.ATPLPolicy; +import io.sentrius.sso.core.utils.JsonUtil; + +public class StringInterpreter implements InputInterpreterIfc { + + @Override + public ATPLPolicy interpret(Map input) throws Exception { + Object policyObj = input.get("policy"); + Object arg1Obj = input.get("arg1"); + + if (policyObj != null) { + return JsonUtil.MAPPER.convertValue(policyObj, ATPLPolicy.class); + } else if (arg1Obj != null) { + return JsonUtil.MAPPER.convertValue(arg1Obj, ATPLPolicy.class); + } + return null; + } +} diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/StringToAtplInterpreter.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/StringToAtplInterpreter.java new file mode 100644 index 00000000..9852f5f4 --- /dev/null +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/interpreters/StringToAtplInterpreter.java @@ -0,0 +1,23 @@ +package io.sentrius.agent.analysis.agents.interpreters; + +import java.util.Map; +import com.fasterxml.jackson.core.JsonParser; +import io.sentrius.agent.analysis.model.TerminalResponse; +import io.sentrius.sso.core.model.verbs.InputInterpreterIfc; +import io.sentrius.sso.core.trust.ATPLPolicy; +import io.sentrius.sso.core.utils.JsonUtil; + +public class StringToAtplInterpreter implements InputInterpreterIfc { + @Override + public ATPLPolicy interpret(Map input) throws Exception { + Object policyObj = input.get("policy"); + Object arg1Obj = input.get("arg1"); + + if (policyObj != null) { + return JsonUtil.MAPPER.convertValue(policyObj, ATPLPolicy.class); + } else if (arg1Obj != null) { + return JsonUtil.MAPPER.convertValue(arg1Obj, ATPLPolicy.class); + } + return null; + } +} diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AgentVerbs.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AgentVerbs.java index 0e85d24d..7f9e6b6f 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AgentVerbs.java +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AgentVerbs.java @@ -10,20 +10,25 @@ import java.util.UUID; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import com.fasterxml.jackson.core.JsonParseException; import com.fasterxml.jackson.core.JsonParser; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import com.google.common.collect.Maps; import com.google.common.collect.Sets; import io.sentrius.agent.analysis.agents.agents.AgentConfig; import io.sentrius.agent.analysis.agents.agents.PromptBuilder; import io.sentrius.agent.analysis.agents.agents.VerbRegistry; +import io.sentrius.agent.analysis.agents.interpreters.AgentContextInterpreter; import io.sentrius.agent.analysis.agents.interpreters.AsessmentListInterpreter; import io.sentrius.agent.analysis.agents.interpreters.ObjectListInterpreter; +import io.sentrius.agent.analysis.agents.interpreters.ObjectNodeInterpreter; +import io.sentrius.agent.analysis.agents.interpreters.StringInterpreter; import io.sentrius.agent.analysis.agents.interpreters.ZtatOutputInterpreter; import io.sentrius.agent.analysis.model.AssessedTerminal; import io.sentrius.agent.analysis.model.Assessment; @@ -32,11 +37,15 @@ import io.sentrius.agent.analysis.model.ZtatAsessment; import io.sentrius.agent.analysis.model.ZtatResponse; import io.sentrius.sso.core.dto.AgentCommunicationDTO; +import io.sentrius.sso.core.dto.AgentRegistrationDTO; import io.sentrius.sso.core.dto.ZtatDTO; +import io.sentrius.sso.core.dto.agents.AgentContextDTO; +import io.sentrius.sso.core.dto.agents.AgentContextRequestDTO; import io.sentrius.sso.core.dto.ztat.AgentExecution; import io.sentrius.sso.core.dto.ztat.AtatRequest; import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO; import io.sentrius.sso.core.exceptions.ZtatException; +import io.sentrius.sso.core.model.verbs.DefaultInterpreter; import io.sentrius.sso.core.model.verbs.Verb; import io.sentrius.sso.core.services.agents.AgentClientService; import io.sentrius.sso.core.services.agents.LLMService; @@ -272,7 +281,8 @@ public String justifyAgent(AgentExecution execution, ZtatRequestDTO ztatRequest, * @throws ZtatException If there is an error during the operation. * @throws IOException If there is an error reading the configuration file. */ - @Verb(name = "assess_data", returnType = ArrayNode.class, description = "Accepts api server data based on the " + + @Verb(name = "assess_api_data", returnType = ArrayNode.class, description = "Accepts api server data based on the" + + " " + "context and seeks" + " to perform the assessment by prompting the LLM. Can be used to assess data or request information from " + "users and/or agents, but not for assessing ztat requests.", @@ -282,7 +292,7 @@ public String justifyAgent(AgentExecution execution, ZtatRequestDTO ztatRequest, public List assessData(AgentExecution execution, List objectList) throws ZtatException, IOException { InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile); if (is == null) { - throw new RuntimeException("agentConfigFile not found on classpath"); + throw new RuntimeException(agentConfigFile + " not found on classpath"); } AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class); @@ -553,4 +563,121 @@ public List analyzeAtatRequests(AgentExecution execution, List messages = new ArrayList<>(); + messages.add( listedEndpoints); + execution.addMessages(Message.builder().role("user").content(queryInput.toString()).build()); + LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build(); + var resp = llmService.askQuestion(execution, chatRequest); + + Response response = JsonUtil.MAPPER.readValue(resp, Response.class); + log.info("Response is {}", resp); + for (Response.Choice choice : response.getChoices()) { + var content = choice.getMessage().getContent(); + if (content.startsWith("```json")) { + content = content.substring(7, content.length() - 3); + } else if (content.startsWith("```")) { + content = content.substring(3, content.length() - 3); + } + log.info("content is {}", content); + if (null != content && !content.isEmpty()) { + try { + + ObjectNode newResponse = JsonUtil.MAPPER.createObjectNode(); + JsonNode node = JsonUtil.MAPPER.readTree(content); + + if (node.isArray()) { + ArrayNode arrayNode = (ArrayNode) node; + newResponse.put("endpoints", arrayNode); + return newResponse; + } else { + log.warn("Expected JSON array but got: {}", node.getNodeType()); + } + + return newResponse; + }catch (JsonParseException e) { + log.error("Failed to parse terminal response: {}", e.getMessage()); + throw e; + } + } + } + return contextNode; + } } \ No newline at end of file diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AtplVerbs.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AtplVerbs.java index 6dd4b758..0e6034b2 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AtplVerbs.java +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AtplVerbs.java @@ -1,27 +1,18 @@ package io.sentrius.agent.analysis.agents.verbs; import java.io.IOException; -import java.net.URLEncoder; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.List; +import java.io.InputStream; import java.util.Map; import com.fasterxml.jackson.databind.node.ArrayNode; -import com.fasterxml.jackson.databind.node.ObjectNode; -import com.google.common.collect.Maps; -import io.sentrius.agent.analysis.agents.interpreters.AsessmentListInterpreter; -import io.sentrius.agent.analysis.agents.interpreters.TerminalListInterpreter; -import io.sentrius.agent.analysis.agents.interpreters.TerminalOutputInterpreter; -import io.sentrius.agent.analysis.model.AssessedTerminal; -import io.sentrius.sso.core.dto.HostSystemDTO; -import io.sentrius.sso.core.dto.ztat.AgentExecution; +import io.sentrius.agent.analysis.agents.interpreters.StringToAtplInterpreter; +import io.sentrius.sso.core.dto.ztat.AtatRequest; import io.sentrius.sso.core.dto.ztat.TokenDTO; -import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO; import io.sentrius.sso.core.exceptions.ZtatException; import io.sentrius.sso.core.model.verbs.DefaultInterpreter; import io.sentrius.sso.core.model.verbs.Verb; import io.sentrius.sso.core.services.agents.LLMService; import io.sentrius.sso.core.services.agents.ZeroTrustClientService; +import io.sentrius.sso.core.trust.ATPLPolicy; import io.sentrius.sso.core.utils.JsonUtil; import lombok.extern.slf4j.Slf4j; import org.springframework.stereotype.Service; @@ -50,18 +41,50 @@ public AtplVerbs(ZeroTrustClientService zeroTrustClientService, LLMService llmSe this.agentVerbs = agentVerbs; } + @Verb(name = "qry_policy_id", description = "Queries by policyId.", + inputInterpreter = StringToAtplInterpreter.class, + outputInterpreter = DefaultInterpreter.class, requiresTokenManagement = true) + public ArrayNode queryPolicyById(TokenDTO token, String policyId) throws ZtatException { + try { + + log.info("policy is : {}", policyId); + String response = zeroTrustClientService.callGetOnApi(token, "/api/v1/policies/" + policyId); + if (response == null) { + throw new RuntimeException("Failed to retrieve terminal list"); + } + log.info("Terminal list response: {}", response); + return (ArrayNode) JsonUtil.MAPPER.readTree(response); + } catch (Exception e) { + throw new RuntimeException("Failed to retrieve terminal list", e); + } + } + + @Verb(name = "get_atpl_schema", description = "Gets Schema. No argument required. Returns JSON Schema.", + inputInterpreter = DefaultInterpreter.class, + outputInterpreter = DefaultInterpreter.class, requiresTokenManagement = true) + public String getAtplSchema(TokenDTO token, Map args) throws ZtatException, IOException { + InputStream schema = getClass().getClassLoader().getResourceAsStream("atpl-schema.json"); + if (schema == null) { + throw new RuntimeException("atpl-schema.json not found on classpath"); + + } + return new String(schema.readAllBytes()); + } + /** * Retrieves a list of currently open terminals. * - * @param args A map of arguments for the operation (currently unused). * @return An `ArrayNode` containing the list of open terminals. * @throws io.sentrius.sso.core.exceptions.ZtatException If there is an error during the operation. */ - @Verb(name = "save_policy", description = "Saves an ATPL policy.", + @Verb(name = "save_policy", description = "Saves an ATPL policy. Accepts ATPL policy in JSON format.", + inputInterpreter = StringToAtplInterpreter.class, outputInterpreter = DefaultInterpreter.class, requiresTokenManagement = true) - public ArrayNode savePolicy(TokenDTO token, Map args) throws ZtatException { + public ArrayNode savePolicy(TokenDTO token, ATPLPolicy policy) throws ZtatException { try { - String response = zeroTrustClientService.callGetOnApi(token, "/ssh/terminal/list/all"); + + log.info("policy is : {}", policy); + String response = zeroTrustClientService.callPostOnApi("/api/v1/policies", policy); if (response == null) { throw new RuntimeException("Failed to retrieve terminal list"); } diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ChatVerbs.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ChatVerbs.java index 9a0d5413..c4890c6e 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ChatVerbs.java +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ChatVerbs.java @@ -7,6 +7,7 @@ import java.nio.file.Paths; import java.util.ArrayList; import java.util.List; +import java.util.ListIterator; import com.fasterxml.jackson.core.JsonParseException; import com.fasterxml.jackson.core.JsonParser; import com.fasterxml.jackson.databind.ObjectMapper; @@ -19,10 +20,13 @@ import io.sentrius.agent.analysis.agents.agents.VerbRegistry; import io.sentrius.agent.analysis.model.TerminalResponse; import io.sentrius.agent.analysis.model.WebSocky; +import io.sentrius.sso.core.dto.agents.AgentContextDTO; +import io.sentrius.sso.core.dto.capabilities.EndpointDescriptor; import io.sentrius.sso.core.dto.ztat.AgentExecution; import io.sentrius.sso.core.exceptions.ZtatException; import io.sentrius.sso.core.model.verbs.Verb; import io.sentrius.sso.core.services.agents.AgentClientService; +import io.sentrius.sso.core.services.agents.AgentExecutionService; import io.sentrius.sso.core.services.agents.LLMService; import io.sentrius.sso.core.services.agents.ZeroTrustClientService; import io.sentrius.sso.core.utils.JsonUtil; @@ -40,9 +44,13 @@ @RequiredArgsConstructor public class ChatVerbs { + private final AgentExecutionService agentExecutionService; @Value("${agent.ai.config}") private String agentConfigFile; + @Value("${agent.ai.context.db.id:none}") + private String agentDatabaseContext; + final ZeroTrustClientService zeroTrustClientService; final LLMService llmService; final VerbRegistry verbRegistry; @@ -71,40 +79,45 @@ public TerminalResponse interpretUserData( throw new RuntimeException("assessor-config.yaml not found on classpath"); } + AgentConfig config = null; String terminalResponse = new String(terminalHelperStream.readAllBytes()); InputStream is = getStream(agentConfigFile); if (is == null) { - throw new RuntimeException(agentConfigFile + " not found on classpath"); + throw new RuntimeException(agentConfigFile + " not found on classpath"); + } + if (agentDatabaseContext != null && !agentDatabaseContext.equals("none")) { + AgentContextDTO agentContext = agentClientService.getAgentContext(execution, + agentDatabaseContext); + config = AgentConfig.builder().description(agentContext.getDescription()) + .context(agentContext.getContext()).build(); + log.info("Agent context loaded: {}", agentContext); + }else { + + config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class); } - AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class); log.info("Agent config loaded: {}", config); PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config); var prompt = promptBuilder.buildPrompt(false ); List messages = new ArrayList<>(); + var context = Message.builder().role("system").content(prompt).build(); + messages.add(context); - messages.add(Message.builder().role("system").content(prompt).build()); - var endpoints = verbRegistry.getEndpoints(); - ArrayNode endpointArray = JsonUtil.MAPPER.createArrayNode(); - for (var verb : endpoints) { - ObjectNode endpoint = JsonUtil.MAPPER.createObjectNode(); - endpoint.put("name", verb.getName()); - - endpoint.put("endpoint", verb.getPath()); - endpointArray.add(endpoint); - } - messages.add(Message.builder().role("system").content("These are a list of available endpoints, " + - "description," + - " their " + - "name:" + endpointArray).build()); + messages.add(Message.builder().role("system").content("You have executed verbs for the previous user " + + "messages. Please generate a user response that summarizes the last message.").build()); + int size = getMessageSize(context); + + var history = getContextWindow(execution.getMessages(), 1024*96 - (size)); + messages.addAll(history); messages.add(Message.builder().role("system").content("Please ensure your nextOperation abides by the " + "following json format and leave it empty if user's request doesn't require explicit use of system " + "operations" + ". Please summarize prior terminal " + "sessions, using " + "terminal output if needed " + - "for clarity of the next LLM request and for the user. Ensure your response meets this json format: " + terminalResponse).build()); + "for clarity of the next LLM request and for the user. Ensure your all future responses meets this " + + "json format: " + terminalResponse).build()); messages.add(Message.builder().role("user").content(userMessage.getContent()).build()); LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build(); var resp = llmService.askQuestion(execution, chatRequest); @@ -128,7 +141,7 @@ public TerminalResponse interpretUserData( return newResponse; }catch (JsonParseException e) { log.error("Failed to parse terminal response: {}", e.getMessage()); - return TerminalResponse.builder().responseForUser(content).terminalSummaryForLLM(lastMessage.getTerminalSummaryForLLM()).build(); + return TerminalResponse.builder().responseForUser(content).build(); } } } @@ -149,30 +162,35 @@ public TerminalResponse interpretUserData( PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config); var prompt = promptBuilder.buildPrompt(false); List messages = new ArrayList<>(); + var context = Message.builder().role("system").content(prompt).build(); + messages.add(context); - messages.add(Message.builder().role("system").content(prompt).build()); - var endpoints = verbRegistry.getEndpoints(); - ArrayNode endpointArray = JsonUtil.MAPPER.createArrayNode(); - for (var verb : endpoints) { - ObjectNode endpoint = JsonUtil.MAPPER.createObjectNode(); - endpoint.put("name", verb.getName()); - - endpoint.put("endpoint", verb.getPath()); - endpointArray.add(endpoint); - } - messages.add(Message.builder().role("system").content("These are a list of available endpoints, " + + /* + var listedEndpoints = Message.builder().role("system").content("These are a list of available endpoints, " + "description," + " their " + - "name:" + endpointArray).build()); - messages.add(Message.builder().role("system").content("Please ensure your nextOperation abide by the " + - "following json format. Please summarize prior terminal sessions, using terminal output if needed " + - "for clarity of the next LLM request and for the user. Ensure your response meets this json format: " + terminalResponse).build()); - messages.add(Message.builder().role("assistant").content("prior response: " + lastMessage.getTerminalSummaryForLLM()).build()); + "name:" + endpointArray).build(); + messages.add(listedEndpoints); + messages.add(Message.builder().role("system").content("You have executed verbs for the previous user " + + "messages. Please generate a user response that summarizes the last message.").build()); + int size = getMessageSize(context) + getMessageSize(listedEndpoints); + + var history = getContextWindow(execution.getMessages(), 1024*96 - (size)); + messages.addAll(history); + + */ + var history = getContextWindow(execution.getMessages(), 1024*96 ); + messages.addAll(history); + +// messages.add(Message.builder().role("assistant").content("prior response: " + lastMessage +// .getTerminalSummaryForLLM()).build()); + + execution.addMessages( userMessage ); messages.add(Message.builder().role("user").content(userMessage.getContent()).build()); LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build(); var resp = llmService.askQuestion(execution, chatRequest); - execution.addMessages( messages ); + Response response = JsonUtil.MAPPER.readValue(resp, Response.class); log.info("Response is {}", resp); for (Response.Choice choice : response.getChoices()) { @@ -184,6 +202,8 @@ public TerminalResponse interpretUserData( } log.info("content is {}", content); if (null != content && !content.isEmpty()) { + + execution.addMessages( choice.getMessage() ); try { var newResponse = JsonUtil.MAPPER.enable(JsonParser.Feature.ALLOW_COMMENTS).readValue( content, @@ -202,6 +222,56 @@ public TerminalResponse interpretUserData( return null; } + public List getContextWindow(List allMessages, int maxContextSize) { + List systemMessages = new ArrayList<>(); + List window = new ArrayList<>(); + int totalSize = 0; + + // First: collect system messages (or other required ones) + for (Message msg : allMessages) { + if ("system".equals(msg.role)) { + systemMessages.add(msg); + totalSize += getMessageSize(msg); + } + } + + // If system messages already exceed max context, return only those + if (totalSize >= maxContextSize) { + return systemMessages; + } + + int remainingSize = maxContextSize - totalSize; + + // Then: collect non-system messages from the end, up to remainingSize + ListIterator iter = allMessages.listIterator(allMessages.size()); + while (iter.hasPrevious()) { + Message msg = iter.previous(); + + if ("system".equals(msg.role)) continue; // already added + + int messageSize = getMessageSize(msg); + if (messageSize > remainingSize) break; + + window.add(0, msg); // prepend + remainingSize -= messageSize; + } + + // Combine system + selected recent messages + List result = new ArrayList<>(); + result.addAll(systemMessages); + result.addAll(window); + + return result; + } + + + private int getMessageSize(Message msg) { + int size = 0; + if (msg.role != null) size += msg.role.length(); + if (msg.content != null) size += msg.content.length(); + if (msg.refusal != null) size += msg.refusal.length(); + return size; + } public TerminalResponse interpret_plan_response( AgentExecution execution, @NonNull WebSocky socketConnection, @@ -230,27 +300,26 @@ public TerminalResponse interpret_plan_response( var prompt = promptBuilder.buildPrompt(false); List messages = new ArrayList<>(); - messages.add(Message.builder().role("system").content(prompt).build()); - var endpoints = verbRegistry.getEndpoints(); - ArrayNode endpointArray = JsonUtil.MAPPER.createArrayNode(); - for (var verb : endpoints) { - ObjectNode endpoint = JsonUtil.MAPPER.createObjectNode(); - endpoint.put("name", verb.getName()); - - endpoint.put("endpoint", verb.getPath()); - endpointArray.add(endpoint); + if (execution.getMessages().isEmpty()) { + var context = Message.builder().role("system").content(prompt).build(); + messages.add(context); + + + } - messages.add(Message.builder().role("system").content("These are a list of available endpoints, " + - "description," + - " their " + - "name:" + endpointArray).build()); + messages.add(Message.builder().role("system").content("You have executed verbs for the previous user " + "messages. Please generate a user response that summarizes the last message.").build()); - if (null != agentVerb ){ - messages.add(Message.builder().role("system").content("You have executed verb: " + agentVerb.getName() + - " with the following description: " + agentVerb.getDescription()).build()); + if (null != agentVerb) { + messages.add( + Message.builder().role("system").content("You have executed verb: " + agentVerb.getName() + + " with the following description: " + agentVerb.getDescription()).build()); } - messages.add(Message.builder().role("assistant").content("prior response: " + lastMessage.getTerminalSummaryForLLM()).build()); + + var history = getContextWindow(execution.getMessages(), 1024*96 ); + messages.addAll(history); + //messages.add(Message.builder().role("assistant").content("prior response: " + lastMessage + // .getTerminalSummaryForLLM()).build()); messages.add(Message.builder().role("assistant").content(planExecutionOutput).build()); LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build(); diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ExampleFactory.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ExampleFactory.java new file mode 100644 index 00000000..418ec15c --- /dev/null +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ExampleFactory.java @@ -0,0 +1,41 @@ +package io.sentrius.agent.analysis.agents.verbs; + +import java.util.List; +import java.util.Map; +import io.sentrius.sso.core.dto.agents.AgentContextDTO; +import io.sentrius.sso.core.trust.ATPLPolicy; +import io.sentrius.sso.core.trust.Capability; +import io.sentrius.sso.core.trust.CapabilitySet; + +public class ExampleFactory { + public static Object createExample(String paramName, Class type) { + if (type.equals(Map.class)) { + return Map.of("key", "value"); + } + if (type.equals(List.class)) { + return List.of(Map.of("key", "value")); + } + if (type.equals(String.class)) { + return "{ \"" + paramName + "\" : Example String value\" }"; + } + if (type.equals(AgentContextDTO.class)) { + return AgentContextDTO.builder().context("This is the context for the agent").description("Agent " + + "description").build(); + } + if (type.equals(ATPLPolicy.class)) { + return ATPLPolicy.builder() + .policyId("policy-001") + .description("Example policy") + .version("v0") + .capabilities( + CapabilitySet.builder().primitives( + List.of( + Capability.builder().description("description").endpoints(List.of( + "endpoint1", "endpoint2")).build() + )).build()) + .build(); + } + // fallback + return Map.of("field", "value"); + } +} diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/api/websocket/ChatWSHandler.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/api/websocket/ChatWSHandler.java index 6ff71c11..d2d4318f 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/analysis/api/websocket/ChatWSHandler.java +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/api/websocket/ChatWSHandler.java @@ -184,8 +184,10 @@ protected void handleTextMessage(WebSocketSession session, TextMessage message) websocketCommunication, userMessage); log.info("Response: {}", response); var newMessage = Session.ChatMessage.newBuilder() - .setMessage(String.format("{\"type\":\"user-message\",\"message\":\"%s\"}", - response.getResponseForUser())) + .setMessage(response.getResponseForUser()/*String.format("{\"type\":\"user-message\"," + + "\"message\":\"%s\"}", + response.getResponseForUser())*/ + ) .setSender("agent") .setChatGroupId("") .setSessionId(Long.parseLong(websocketCommunication.getSessionId())) @@ -206,10 +208,11 @@ protected void handleTextMessage(WebSocketSession session, TextMessage message) var lastVerbResponse = websocketCommunication.getVerbResponses().stream().reduce((prev, next) -> next) .orElse(null); + var arguments = response.getArguments(); var executionResponse = verbRegistry.execute( chatAgent.getAgentExecution(), lastVerbResponse, - response.getNextOperation(), Maps.newHashMap() + response.getNextOperation(), arguments ); var nextResponse = chatVerbs.interpret_plan_response( diff --git a/ai-agent/src/main/java/io/sentrius/agent/analysis/model/TerminalResponse.java b/ai-agent/src/main/java/io/sentrius/agent/analysis/model/TerminalResponse.java index 02655a82..5bcacf37 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/analysis/model/TerminalResponse.java +++ b/ai-agent/src/main/java/io/sentrius/agent/analysis/model/TerminalResponse.java @@ -2,6 +2,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; import io.sentrius.sso.core.dto.HostSystemDTO; import io.sentrius.sso.genai.Message; import lombok.AllArgsConstructor; @@ -23,4 +24,6 @@ public class TerminalResponse { String nextOperation; String terminalSummaryForLLM; String responseForUser; + @Builder.Default + public Map arguments = Map.of(); } diff --git a/ai-agent/src/main/java/io/sentrius/agent/config/SecurityConfig.java b/ai-agent/src/main/java/io/sentrius/agent/config/SecurityConfig.java index 398de64d..ff352101 100644 --- a/ai-agent/src/main/java/io/sentrius/agent/config/SecurityConfig.java +++ b/ai-agent/src/main/java/io/sentrius/agent/config/SecurityConfig.java @@ -66,9 +66,7 @@ public JwtAuthenticationConverter jwtAuthenticationConverterForKeycloak() { log.info("**** Initializing JwtAuthenticationConverter"); converter.setJwtGrantedAuthoritiesConverter(jwt -> { - log.info("**** Jwt Authentication Converter invoked"); Collection authorities = new JwtGrantedAuthoritiesConverter().convert(jwt); - log.info("JWT Claims: {}", jwt.getClaims()); String userId = jwt.getClaimAsString("sub"); String username = jwt.getClaimAsString("preferred_username"); diff --git a/ai-agent/src/main/resources/atpl-schema.json b/ai-agent/src/main/resources/atpl-schema.json new file mode 100644 index 00000000..14839924 --- /dev/null +++ b/ai-agent/src/main/resources/atpl-schema.json @@ -0,0 +1,199 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://sentrius.io/schemas/atpl.schema.json", + "title": "Agent Trust Policy Language (ATPL)", + "type": "object", + "required": ["version", "policy_id", "capabilities"], + "properties": { + "version": { + "type": "string", + "pattern": "^v[0-9]+$" + }, + "policy_id": { + "type": "string" + }, + "description": { + "type": "string" + }, + "match": { + "type": "object", + "properties": { + "agent_tags": { + "type": "array", + "items": { "type": "string" } + } + } + }, + "identity": { + "type": "object", + "properties": { + "issuer": { "type": "string" }, + "subject_prefix": { "type": "string" }, + "mfa_required": { "type": "boolean" }, + "certificate_authority": { "type": "string" } + } + }, + "provenance": { + "type": "object", + "properties": { + "source": { "type": "string" }, + "signature_required": { "type": "boolean" }, + "approved_committers": { + "type": "array", + "items": { "type": "string" } + } + } + }, + "runtime": { + "type": "object", + "properties": { + "allow_drift": { "type": "boolean" }, + "enclave_required": { "type": "boolean" }, + "attestation_type": { "type": "string" }, + "verified_at_boot": { "type": "boolean" } + } + }, + "behavior": { + "type": "object", + "properties": { + "minimum_positive_runs": { "type": "integer" }, + "max_incidents": { "type": "integer" }, + "incident_types": { + "type": "object", + "properties": { + "denylist": { + "type": "array", + "items": { "type": "string" } + } + } + } + } + }, + "trust_score": { + "type": "object", + "properties": { + "minimum": { "type": "integer" }, + "weightings": { + "type": "object", + "properties": { + "identity": { "type": "number" }, + "provenance": { "type": "number" }, + "runtime": { "type": "number" }, + "behavior": { "type": "number" } + } + } + } + }, + "capabilities": { + "type": "object", + "required": ["primitives", "composed"], + "properties": { + "primitives": { + "type": "array", + "items": { + "type": "object", + "required": ["id", "description"], + "properties": { + "id": { "type": "string" }, + "description": { "type": "string" }, + "endpoints": { + "type": "array", + "items": { "type": "string" } + }, + "commands": { + "type": "array", + "items": { "type": "string" } + }, + "subcommands": { + "type": "array", + "items": { "type": "string" } + }, + "activities": { + "type": "array", + "items": { "type": "string" } + }, + "tags": { + "type": "array", + "items": { "type": "string" } + } + } + } + }, + "composed": { + "type": "array", + "items": { + "type": "object", + "required": ["id", "includes"], + "properties": { + "id": { "type": "string" }, + "includes": { + "type": "array", + "items": { "type": "string" } + }, + "tags": { + "type": "array", + "items": { "type": "string" } + }, + "endpoints": { + "type": "array", + "items": { "type": "string" } + }, + "commands": { + "type": "array", + "items": { "type": "string" } + }, + "subcommands": { + "type": "array", + "items": { "type": "string" } + }, + "activities": { + "type": "array", + "items": { "type": "string" } + } + } + } + } + } + } + , + "actions": { + "type": "object", + "properties": { + "on_failure": { + "type": "string", + "enum": ["deny", "log", "alert"] + }, + "on_success": { + "type": "string", + "enum": ["allow", "log"] + }, + "on_marginal": { + "oneOf": [ + {"type": "string", "enum": ["require_ztat", "log"]}, + { + "type": "object", + "properties": { + "action": {"type": "string", "enum": ["require_ztat"]}, + "ztat_provider": {"type": "string"} + }, + "required": ["action"] + } + ] + } + } + }, + "ztat": { + "type": "object", + "properties": { + "provider": {"type": "string"}, + "ttl": {"type": "string"}, + "approved_issuers": { + "type": "array", + "items": {"type": "string"} + }, + "key_binding": {"type": "string"}, + "approval_required": {"type": "boolean"} + } + } + } +} diff --git a/ai-agent/src/main/resources/chat-helper.json b/ai-agent/src/main/resources/chat-helper.json index 7bc969b9..bffb7174 100644 --- a/ai-agent/src/main/resources/chat-helper.json +++ b/ai-agent/src/main/resources/chat-helper.json @@ -1,6 +1,10 @@ { "previousOperation": "", "nextOperation": "", + "arguments": { + "argumentname": "", + "argumentname2": "" + }, "terminalSummaryForLLM": "", "responseForUser": "" } \ No newline at end of file diff --git a/ai-agent/src/main/resources/terminal-helper.json b/ai-agent/src/main/resources/terminal-helper.json index 7bc969b9..02be454b 100644 --- a/ai-agent/src/main/resources/terminal-helper.json +++ b/ai-agent/src/main/resources/terminal-helper.json @@ -1,6 +1,10 @@ { "previousOperation": "", "nextOperation": "", + "arguments": { + "argumentname": "", + "argumentname2": "" + }, "terminalSummaryForLLM": "", - "responseForUser": "" + "responseForUser": "" } \ No newline at end of file diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java b/api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java index 40a303e3..ab573e2d 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java +++ b/api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java @@ -15,11 +15,16 @@ import java.util.concurrent.ExecutionException; import com.fasterxml.jackson.core.JsonProcessingException; import io.sentrius.sso.config.ApiPaths; +import io.sentrius.sso.config.AppConfig; import io.sentrius.sso.core.annotations.LimitAccess; import io.sentrius.sso.core.config.SystemOptions; import io.sentrius.sso.core.controllers.BaseController; import io.sentrius.sso.core.dto.AgentCommunicationDTO; +import io.sentrius.sso.core.dto.AgentDTO; import io.sentrius.sso.core.dto.AgentHeartbeatDTO; +import io.sentrius.sso.core.dto.agents.AgentContextDTO; +import io.sentrius.sso.core.dto.agents.AgentContextRequestDTO; +import io.sentrius.sso.core.exceptions.ZtatException; import io.sentrius.sso.core.model.chat.AgentCommunication; import io.sentrius.sso.core.model.security.enums.IdentityType; import io.sentrius.sso.core.model.security.UserType; @@ -32,6 +37,8 @@ import io.sentrius.sso.core.services.ATPLPolicyService; import io.sentrius.sso.core.services.ErrorOutputService; import io.sentrius.sso.core.services.UserService; +import io.sentrius.sso.core.services.agents.AgentClientService; +import io.sentrius.sso.core.services.agents.AgentContextService; import io.sentrius.sso.core.services.agents.AgentService; import io.sentrius.sso.core.services.auditing.AuditService; import io.sentrius.sso.core.services.security.CryptoService; @@ -54,6 +61,7 @@ import org.springframework.format.annotation.DateTimeFormat; import org.springframework.http.ResponseEntity; import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.PathVariable; import org.springframework.web.bind.annotation.PostMapping; import org.springframework.web.bind.annotation.RequestBody; import org.springframework.web.bind.annotation.RequestHeader; @@ -75,7 +83,10 @@ public class AgentApiController extends BaseController { final ZeroTrustRequestService ztrService; final AgentService agentService; final ProvenanceKafkaProducer provenanceKafkaProducer; - private final ZeroTrustRequestService ztatRequestService; + final ZeroTrustRequestService ztatRequestService; + final AgentContextService agentContextService; + final AgentClientService agentClientService; + final AppConfig appConfig; public AgentApiController( UserService userService, @@ -85,7 +96,8 @@ public AgentApiController( CryptoService cryptoService, SessionTrackingService sessionTrackingService, KeycloakService keycloakService, ATPLPolicyService atplPolicyService, ZeroTrustAccessTokenService ztatService, ZeroTrustRequestService ztrService, AgentService agentService, - ProvenanceKafkaProducer provenanceKafkaProducer, ZeroTrustRequestService ztatRequestService + ProvenanceKafkaProducer provenanceKafkaProducer, ZeroTrustRequestService ztatRequestService, + AgentContextService agentContextService, AgentClientService agentClientService, AppConfig appConfig ) { super(userService, systemOptions, errorOutputService); this.auditService = auditService; @@ -98,6 +110,9 @@ public AgentApiController( this.agentService = agentService; this.provenanceKafkaProducer = provenanceKafkaProducer; this.ztatRequestService = ztatRequestService; + this.agentContextService = agentContextService; + this.agentClientService = agentClientService; + this.appConfig = appConfig; } public SessionLog createSession(@RequestParam String username, @RequestParam String ipAddress) { @@ -282,7 +297,7 @@ public ResponseEntity createAgentChatRequest( @GetMapping("/list") @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) - public ResponseEntity listAgents(HttpServletRequest request, HttpServletResponse response){ + public ResponseEntity listAgents(HttpServletRequest request, HttpServletResponse response) throws ZtatException { var operatingUser = getOperatingUser(request, response ); if (null == operatingUser) { log.warn("No operating user found"); @@ -290,7 +305,24 @@ public ResponseEntity listAgents(HttpServletRequest request, HttpServletRespo } log.info("Received list request from user: {} {}", operatingUser.getUsername(), operatingUser); var agents = agentService.getAllAgents(true); - return ResponseEntity.ok(agents); + + List prunedAgentList = agents.stream().filter(agent -> { + try { + String podResponse = + agentClientService.getAgentPodStatus(appConfig.getSentriusLauncherService(), agent.getAgentName()); + if (podResponse != null && (podResponse.equalsIgnoreCase("running") || podResponse.equalsIgnoreCase("pending"))){ + return true; + } else { + log.info("Agent {} is not running or pending, removing from list. Status is {}", + agent.getAgentName(), podResponse); + } + } catch (ZtatException ignored) { + + } + return false; + } + ).toList(); + return ResponseEntity.ok(prunedAgentList); } @@ -754,4 +786,38 @@ private boolean validateUser(User requestor, User operatingUser, AgentCommunicat return canSend; } + @GetMapping("/context/{contextId}") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public ResponseEntity getContext( + HttpServletRequest request, + HttpServletResponse response, + @PathVariable("contextId") String contextId){ + var databaseContext = agentContextService.getContextOrThrow(UUID.fromString(contextId)); + return ResponseEntity.ok(AgentContextDTO.builder() + .id(databaseContext.getId()) + .name(databaseContext.getName()) + .description(databaseContext.getDescription()) + .context(databaseContext.getContext()) + .createdAt(databaseContext.getCreatedAt()) + .updatedAt(databaseContext.getUpdatedAt()) + .build()); + } + + @PostMapping("/context") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public ResponseEntity createContext( + HttpServletRequest request, + HttpServletResponse response, + AgentContextRequestDTO dtoRequest){ + var databaseContext = agentContextService.create(dtoRequest); + return ResponseEntity.ok(AgentContextDTO.builder() + .id(databaseContext.getId()) + .name(databaseContext.getName()) + .description(databaseContext.getDescription()) + .context(databaseContext.getContext()) + .createdAt(databaseContext.getCreatedAt()) + .updatedAt(databaseContext.getUpdatedAt()) + .build()); + } + } diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/AgentBootstrapController.java b/api/src/main/java/io/sentrius/sso/controllers/api/AgentBootstrapController.java index f8ecdc3f..147f2ac5 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/api/AgentBootstrapController.java +++ b/api/src/main/java/io/sentrius/sso/controllers/api/AgentBootstrapController.java @@ -40,6 +40,7 @@ import org.springframework.beans.factory.annotation.Value; import org.springframework.http.ResponseEntity; import org.springframework.web.bind.annotation.DeleteMapping; +import org.springframework.web.bind.annotation.GetMapping; import org.springframework.web.bind.annotation.PostMapping; import org.springframework.web.bind.annotation.RequestBody; import org.springframework.web.bind.annotation.RequestHeader; @@ -62,6 +63,7 @@ public class AgentBootstrapController extends BaseController { private final ZeroTrustClientService zeroTrustClientService; private final ObjectMapper yamlMapper = new ObjectMapper(new YAMLFactory()); final AppConfig appConfig; + private final AgentClientService agentClientService; public AgentBootstrapController( @@ -72,7 +74,8 @@ public AgentBootstrapController( CryptoService cryptoService, SessionTrackingService sessionTrackingService, KeycloakService keycloakService, ATPLPolicyService atplPolicyService, ZeroTrustAccessTokenService ztatService, ZeroTrustRequestService ztrService, AgentService agentService, - ZeroTrustClientService zeroTrustClientService, AppConfig appConfig + ZeroTrustClientService zeroTrustClientService, AppConfig appConfig, + AgentClientService agentClientService ) { super(userService, systemOptions, errorOutputService); this.auditService = auditService; @@ -85,6 +88,7 @@ public AgentBootstrapController( this.agentService = agentService; this.zeroTrustClientService = zeroTrustClientService; this.appConfig = appConfig; + this.agentClientService = agentClientService; } @@ -168,7 +172,26 @@ public ResponseEntity launchPod( @RequestBody AgentRegistrationDTO registrationDTO, HttpServletRequest request, HttpServletResponse response ) throws GeneralSecurityException, IOException, ZtatException { + try{ + log.info("Launching agent pod with ID: {}", registrationDTO.getAgentName()); + var status = getAgentStatus( registrationDTO.getAgentName(), request, response); + if ( status != null ) { + var body = status.getBody(); + if (body != null) { + + if (body.contains("Running") || body.contains("Pending")) { + log.info("Agent {} is already running or pending", registrationDTO.getAgentName()); + return ResponseEntity.ok("{\"status\": \"already exists\"}"); + } else { + log.warn("Agent {} is not running, attempting to launch again", registrationDTO.getAgentName()); + } + } + } + } catch (Exception e) { + log.error("Error getting agent status", e); + + } var operatingUser = getOperatingUser(request, response ); zeroTrustClientService.callAuthenticatedPostOnApi(appConfig.getSentriusLauncherService(), "agent/launcher/create", registrationDTO); @@ -192,6 +215,19 @@ public ResponseEntity deletePod( } + @GetMapping("/launcher/status") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public ResponseEntity getAgentStatus( + @RequestParam("agentId") String agentId, HttpServletRequest request, HttpServletResponse response + ) throws GeneralSecurityException, IOException, ZtatException { + + + var operatingUser = getOperatingUser(request, response ); + String podResponse = + agentClientService.getAgentPodStatus(appConfig.getSentriusLauncherService(), agentId); + // bootstrap with a default policy + return ResponseEntity.ok("{\"status\": \"" + podResponse + "\"}"); + } diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/HostApiController.java b/api/src/main/java/io/sentrius/sso/controllers/api/HostApiController.java index 8306c824..de7b96a5 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/api/HostApiController.java +++ b/api/src/main/java/io/sentrius/sso/controllers/api/HostApiController.java @@ -196,7 +196,7 @@ public ResponseEntity connectSSHServer(HttpServletRequest request, H if (enclaveId == null || hostId == null) { return ResponseEntity.badRequest().build(); } - if (systemOptions.getSshEnabled() == false){ + if (systemOptions.getLockdownEnabled() == true){ node.put("sessionId",""); node.put("errorToUser","SSH is disabled"); return ResponseEntity.ok(node); diff --git a/api/src/main/java/io/sentrius/sso/controllers/api/SystemApiController.java b/api/src/main/java/io/sentrius/sso/controllers/api/SystemApiController.java index ccedfb67..26721f4e 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/api/SystemApiController.java +++ b/api/src/main/java/io/sentrius/sso/controllers/api/SystemApiController.java @@ -73,20 +73,20 @@ protected SystemApiController( this.configurationApplicationTask = configurationApplicationTask; } - @GetMapping("/settings/sshEnabled") - public ResponseEntity getSSHEnabled() { + @GetMapping("/settings/lockdownEnabled") + public ResponseEntity getLockdownEnabled() { ObjectNode node = JsonUtil.MAPPER.createObjectNode(); - node.put("sshEnabled", systemOptions.getSshEnabled()); + node.put("lockdownEnabled", systemOptions.getLockdownEnabled()); return ResponseEntity.ok(node); } - @PutMapping("/settings/ssh/toggle") + @PutMapping("/settings/lockdown/toggle") @LimitAccess(applicationAccess ={ ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) - public ResponseEntity toggleSSHEnabled() { - log.info("Toggling SSH enabled"); + public ResponseEntity toggleLockdown() { + log.info("Toggling Lockdown enabled"); ObjectNode node = JsonUtil.MAPPER.createObjectNode(); - systemOptions.setValue("sshEnabled", !systemOptions.getSshEnabled()); - node.put("sshEnabled", systemOptions.getSshEnabled()); + systemOptions.setValue("lockdownEnabled", !systemOptions.getLockdownEnabled()); + node.put("lockdownEnabled", systemOptions.getLockdownEnabled()); return ResponseEntity.ok(node); } diff --git a/api/src/main/java/io/sentrius/sso/controllers/view/AgentController.java b/api/src/main/java/io/sentrius/sso/controllers/view/AgentController.java index 58bcf2b9..3670cdd1 100644 --- a/api/src/main/java/io/sentrius/sso/controllers/view/AgentController.java +++ b/api/src/main/java/io/sentrius/sso/controllers/view/AgentController.java @@ -48,6 +48,12 @@ public String listAgents(Model m) { return "sso/agents/list_agents"; } + @GetMapping("/design/chat") + @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) + public String designAgent(Model m) { + return "sso/agents/design_chat"; + } + @GetMapping("/connections") @LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION}) public String listConnections(Model m, @RequestParam("agentId") String agentId) throws GeneralSecurityException { @@ -59,4 +65,5 @@ public String listConnections(Model m, @RequestParam("agentId") String agentId) return "sso/agents/agent_comms"; } + } diff --git a/api/src/main/resources/db/migration/V18_agent_launch.sql b/api/src/main/resources/db/migration/V18_agent_launch.sql new file mode 100644 index 00000000..63694ba6 --- /dev/null +++ b/api/src/main/resources/db/migration/V18_agent_launch.sql @@ -0,0 +1,17 @@ +CREATE TABLE agent_contexts ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + name TEXT NOT NULL, + description TEXT, + context TEXT NOT NULL, -- YAML or any other string format + created_at TIMESTAMP WITH TIME ZONE DEFAULT now(), + updated_at TIMESTAMP WITH TIME ZONE DEFAULT now() +); + +CREATE TABLE agent_launches ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + agent_id TEXT NOT NULL, -- e.g., name or UUID of the launched agent + context_id UUID NOT NULL REFERENCES agent_contexts(id), + launched_by TEXT, -- who or what initiated it + launch_parameters TEXT, -- optional overrides or launch args + created_at TIMESTAMP WITH TIME ZONE DEFAULT now() +); \ No newline at end of file diff --git a/api/src/main/resources/static/js/add_system.js b/api/src/main/resources/static/js/add_system.js index 3752b900..0ae07c5e 100644 --- a/api/src/main/resources/static/js/add_system.js +++ b/api/src/main/resources/static/js/add_system.js @@ -9,14 +9,14 @@ document.addEventListener('DOMContentLoaded', function () { if (disableSSHButton) { - fetch(`/api/v1/system/settings/sshEnabled`) + fetch(`/api/v1/system/settings/lockdownEnabled`) .then(response => response.json()) .then(data => { - if (data.sshEnabled) { - disableSSHButton.innerText = 'Disable SSH'; + if (data.lockdownEnabled) { + disableSSHButton.innerText = 'LockDown Systems'; } else { - disableSSHButton.innerText = 'Enable SSH'; + disableSSHButton.innerText = 'Re-enable Systems'; } }) .catch(error => { @@ -25,7 +25,7 @@ document.addEventListener('DOMContentLoaded', function () { document.getElementById('disable-ssh-button').addEventListener('click', function(event) { event.preventDefault(); // Prevent the default anchor behavior const csrfToken = document.getElementById('csrf-token').value; // Get CSRF token value - fetch('/api/v1/system/settings/ssh/toggle', { + fetch('/api/v1/system/settings/lockdown/toggle', { method: 'PUT', // Specify PUT request headers: { 'Content-Type': 'application/json', // Optional, adjust based on your API @@ -34,10 +34,10 @@ document.addEventListener('DOMContentLoaded', function () { }).then(response => response.json()) .then(data => { if (data.sshEnabled) { - disableSSHButton.innerText = 'Disable SSH'; + disableSSHButton.innerText = 'LockDown Systems'; } else { - disableSSHButton.innerText = 'Enable SSH'; + disableSSHButton.innerText = 'Re-enable Systems'; } }) .catch(error => { diff --git a/api/src/main/resources/templates/sso/agents/design_chat.html b/api/src/main/resources/templates/sso/agents/design_chat.html new file mode 100644 index 00000000..789b180b --- /dev/null +++ b/api/src/main/resources/templates/sso/agents/design_chat.html @@ -0,0 +1,692 @@ + + + + + [[${systemOptions.systemLogoName}]] - Agent Designer + + + + + +
+
+ +
+
+
+
+ +
+
+

Agent Designer Assistant

+

Let me help you create an agent

+
+ +
+
+ Welcome! I'm here to help you create an agent. + What kind of agent are you looking to configure? +
+
+ +
+
+ + + + +
+ +
+ + +
+
+ +
+
Assistant is thinking... +
+
+ +
+ + + +
+
+
+
+ + + + \ No newline at end of file diff --git a/api/src/main/resources/templates/sso/dashboard.html b/api/src/main/resources/templates/sso/dashboard.html index 60006ff0..adc32699 100755 --- a/api/src/main/resources/templates/sso/dashboard.html +++ b/api/src/main/resources/templates/sso/dashboard.html @@ -351,168 +351,168 @@ }); -