Skip to content

Commit 6941aa7

Browse files
phrockerphrocker
committed
proxy Chat
1 parent e5080e6 commit 6941aa7

File tree

22 files changed

+341
-30
lines changed

22 files changed

+341
-30
lines changed

‎.local.env‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
SENTRIUS_VERSION=1.1.81
2-
SENTRIUS_SSH_VERSION=1.1.17
3-
SENTRIUS_KEYCLOAK_VERSION=1.1.23
4-
SENTRIUS_AGENT_VERSION=1.1.17
5-
SENTRIUS_AI_AGENT_VERSION=1.1.32
6-
LLMPROXY_VERSION=1.0.17
7-
LAUNCHER_VERSION=1.0.22
1+
SENTRIUS_VERSION=1.1.92
2+
SENTRIUS_SSH_VERSION=1.1.18
3+
SENTRIUS_KEYCLOAK_VERSION=1.1.25
4+
SENTRIUS_AGENT_VERSION=1.1.18
5+
SENTRIUS_AI_AGENT_VERSION=1.1.33
6+
LLMPROXY_VERSION=1.0.18
7+
LAUNCHER_VERSION=1.0.27

‎.local.env.bak‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
SENTRIUS_VERSION=1.1.81
2-
SENTRIUS_SSH_VERSION=1.1.17
3-
SENTRIUS_KEYCLOAK_VERSION=1.1.22
4-
SENTRIUS_AGENT_VERSION=1.1.17
5-
SENTRIUS_AI_AGENT_VERSION=1.1.32
6-
LLMPROXY_VERSION=1.0.17
7-
LAUNCHER_VERSION=1.0.22
1+
SENTRIUS_VERSION=1.1.92
2+
SENTRIUS_SSH_VERSION=1.1.18
3+
SENTRIUS_KEYCLOAK_VERSION=1.1.25
4+
SENTRIUS_AGENT_VERSION=1.1.18
5+
SENTRIUS_AI_AGENT_VERSION=1.1.33
6+
LLMPROXY_VERSION=1.0.18
7+
LAUNCHER_VERSION=1.0.26

‎agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,19 @@ public class PodLauncherService {
2828
@Value("${sentrius.agent.registry.version}")
2929
private String agentVersion;
3030

31+
@Value("${sentrius.agent.callback.format.url:http://sentrius-agent-%s.%s.svc.cluster.local:8090}")
32+
private String callbackFormatUrl;
33+
3134
public PodLauncherService() throws IOException {
3235
ApiClient client = Config.defaultClient(); // in-cluster or kubeconfig
3336
this.coreV1Api = new CoreV1Api(client);
3437
}
3538

39+
private String buildAgentCallbackUrl(String agentId) {
40+
return String.format(callbackFormatUrl, agentId, agentNamespace);
41+
}
42+
43+
3644
public V1Pod launchAgentPod(String agentId, String callbackUrl) throws Exception {
3745
if (agentRegistry != null ) {
3846
if ("local".equalsIgnoreCase(agentRegistry)) {
@@ -42,6 +50,8 @@ public V1Pod launchAgentPod(String agentId, String callbackUrl) throws Exception
4250
}
4351
}
4452

53+
var constructedCallbackUrl = buildAgentCallbackUrl(agentId);
54+
4555
String image = String.format("%ssentrius-launchable-agent:%s", agentRegistry, agentVersion);
4656

4757
log.info("Launching agent pod with ID: {}, Image: {}, Callback URL: {}", agentId, image, callbackUrl);
@@ -56,7 +66,9 @@ public V1Pod launchAgentPod(String agentId, String callbackUrl) throws Exception
5666
.imagePullPolicy("IfNotPresent")
5767

5868
.args(List.of("--spring.config.location=file:/config/agent.properties",
59-
"--agent.namePrefix=" + agentId, "--agent.ai.config=/config/chat-helper.yaml", "--agent.listen.websocket=true"))
69+
"--agent.namePrefix=" + agentId, "--agent.ai.config=/config/chat-helper.yaml", "--agent.listen.websocket=true",
70+
"--agent.callback.url=" + constructedCallbackUrl
71+
))
6072
.resources(new V1ResourceRequirements()
6173
.limits(Map.of(
6274
"cpu", Quantity.fromString("500m"),

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@
33
import java.net.URLDecoder;
44
import java.nio.charset.StandardCharsets;
55
import java.security.GeneralSecurityException;
6+
import java.security.NoSuchAlgorithmException;
67
import java.sql.SQLException;
78
import java.time.LocalDateTime;
89
import java.util.HashSet;
910
import java.util.List;
1011
import java.util.Map;
12+
import java.util.Optional;
1113
import java.util.Set;
1214
import java.util.UUID;
1315
import java.util.concurrent.ExecutionException;
@@ -40,6 +42,7 @@
4042
import io.sentrius.sso.provenance.kafka.ProvenanceKafkaProducer;
4143
import jakarta.servlet.http.HttpServletRequest;
4244
import jakarta.servlet.http.HttpServletResponse;
45+
import jakarta.transaction.Transactional;
4346
import lombok.extern.slf4j.Slf4j;
4447
import org.apache.http.HttpStatus;
4548
import org.jetbrains.annotations.NotNull;
@@ -172,7 +175,7 @@ public ResponseEntity<?> requestRegistration(
172175

173176
// Approve the request if the agent has an active policy ( and it is known and allowed ).
174177
if (atplPolicyService.getPolicy(operatingUser).isPresent()) {
175-
var admin = userService.getUser(UserType.createSystemAdmin().getId());
178+
var admin = createOrGetSystemAdmin();
176179
var approval = ztatService.approveOpsAccessToken(ztatRequest, admin);
177180

178181
return ResponseEntity.ok(Map.of("ztat_token", approval.getToken().toString(), "communication_id",communicationId ));
@@ -186,6 +189,21 @@ public ResponseEntity<?> requestRegistration(
186189

187190

188191

192+
}
193+
194+
@Transactional
195+
protected synchronized User createOrGetSystemAdmin() throws NoSuchAlgorithmException {
196+
var admin = userService.getUserByUsername("SYSTEM");
197+
if (null == admin){
198+
var systemAdmin = User.builder()
199+
.username("SYSTEM")
200+
.name("System Admin")
201+
.userId("SYSTEM")
202+
.emailAddress("email").password( userService.encodePassword(UUID.randomUUID().toString())).authorizationType(UserType.createSystemAdmin()).identityType(IdentityType.NON_PERSON_ENTITY);
203+
return userService.save(systemAdmin.build());
204+
}
205+
return admin;
206+
189207
}
190208

191209
@PostMapping("/provenance/submit")

‎api/src/main/java/io/sentrius/sso/controllers/api/EnclaveApiController.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,8 +75,8 @@ public ResponseEntity<HostGroupDTO> setAssignments(HttpServletRequest request, H
7575
List<User> newUserList = new ArrayList<>();
7676
for(var userId : (List<String>) payload.get("userIds")) {
7777
var u = userService.getUser(Long.valueOf(userId));
78-
if (null != u) {
79-
newUserList.add(u);
78+
if (u.isPresent()) {
79+
newUserList.add(u.get());
8080
}
8181
}
8282
hg.setUsers(newUserList);

‎api/src/main/java/io/sentrius/sso/controllers/api/UserApiController.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,9 +142,19 @@ public String deleteUser(@RequestParam("userId") String userId, @RequestParam(re
142142
log.info("Deleting non-person entity user with id: {}", userId);
143143
String userIdStr = cryptoService.decrypt(userId);
144144
var usr = userService.getUserByUserid(userIdStr);
145+
if (usr.getId() < 0) {
146+
log.info("User with id {} is a system user and cannot be deleted", usr.getId());
147+
return "redirect:/sso/v1/users/list?message=" + MessagingUtil.getMessageId(MessagingUtil.UNEXPECTED_ERROR);
148+
149+
}
145150
userService.deleteUser(usr.getId());
146151
} else {
147152
Long id = Long.parseLong(cryptoService.decrypt(userId));
153+
if (id < 0) {
154+
log.info("User with id {} is a system user and cannot be deleted", id);
155+
return "redirect:/sso/v1/users/list?message=" +
156+
MessagingUtil.getMessageId(MessagingUtil.UNEXPECTED_ERROR);
157+
}
148158
userService.deleteUser(id);
149159
}
150160
return "redirect:/sso/v1/users/list?message=" + MessagingUtil.getMessageId(MessagingUtil.USER_DELETE_SUCCESS);

‎api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
package io.sentrius.sso.locator;
2+
3+
4+
import org.springframework.beans.factory.annotation.Value;
5+
import org.springframework.stereotype.Component;
6+
7+
import java.net.URI;
8+
9+
@Component
10+
public class KubernetesAgentLocator {
11+
12+
@Value("${sentrius.agent.namespace}")
13+
private String agentNamespace;
14+
15+
@Value("${sentrius.agent.port:8080}")
16+
private int agentPort;
17+
18+
public URI resolveWebSocketUri(String agentId) {
19+
// DNS: sentrius-agent-[ID].[namespace].svc.cluster.local
20+
String fqdn = String.format("ws://sentrius-agent-%s.%s.svc.cluster.local:%d/ws",
21+
agentId, agentNamespace, agentPort);
22+
return URI.create(fqdn);
23+
}
24+
}

‎api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java‎

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,8 @@ public List<SideEffect> initialize(InstallConfiguration installConfiguration, bo
189189

190190
// first we create the admin user, then the user types followed by all users
191191
sideEffects.addAll(createAdminUser(installConfiguration, action));
192+
sideEffects.addAll(createSystemAdmin(installConfiguration, action));
193+
192194

193195
createSystemUser(installConfiguration);
194196

@@ -681,7 +683,7 @@ protected List<User> createUsers(
681683
}
682684
}
683685
if (action){
684-
user = userService.getUser(user.getId());
686+
user = userService.getUser(user.getId()).orElseThrow();
685687
var definition = userDTO.getAtlpDefinition();
686688
if (null != definition && !definition.isEmpty()) {
687689
Optional<ATPLPolicyEntity> policy = policyList.stream()
@@ -809,7 +811,7 @@ protected List<User> createNPEs(
809811
}
810812
}
811813
if (action){
812-
user = userService.getUser(user.getId());
814+
user = userService.getUser(user.getId()).orElseThrow();
813815
var definition = userDTO.getAtlpDefinition();
814816
if (null != definition && !definition.isEmpty()) {
815817
Optional<ATPLPolicyEntity> policy = policyList.stream()
@@ -880,6 +882,50 @@ protected List<SideEffect> createAdminUser(InstallConfiguration installConfigura
880882
return sideEffects;
881883
}
882884

885+
@Transactional
886+
public List<SideEffect> createSystemAdmin(InstallConfiguration installConfiguration, boolean action) throws NoSuchAlgorithmException {
887+
888+
var user = installConfiguration.getSystemUser();
889+
890+
if (null == user) {
891+
throw new IllegalStateException("Admin user not found in configuration");
892+
}
893+
List<SideEffect> sideEffects = new ArrayList<>();
894+
userService.findByUsername("SYSTEM").ifPresentOrElse(
895+
user1 -> {
896+
// ignore
897+
},
898+
() -> {
899+
sideEffects.add(SideEffect.builder().sideEffectDescription("Creating admin user " + user.getUsername()).type(
900+
SideEffectType.UPDATE_DATABASE).asset("Users").build());
901+
if (action) {
902+
try {
903+
user.setUserId("SYSTEM");
904+
user.setPassword(userService.encodePassword(UUID.randomUUID().toString()));
905+
user.setAuthorizationType(UserType.createSystemAdmin().toDTO());
906+
user.setIdentityType(IdentityType.NON_PERSON_ENTITY.toString());
907+
908+
var type =
909+
userService.getUserType(UserType.createSystemAdmin());
910+
if (type.isEmpty()){
911+
type = Optional.of( userService.saveUserType(UserType.createSystemAdmin()) );
912+
}
913+
914+
userService.addUscer(User.from(user, type.get()));
915+
} catch (NoSuchAlgorithmException e) {
916+
throw new RuntimeException(e);
917+
}
918+
919+
// insert default admin user
920+
921+
}
922+
}
923+
);
924+
925+
926+
return sideEffects;
927+
}
928+
883929
@Transactional
884930
protected void createSystemUser(InstallConfiguration connection) throws NoSuchAlgorithmException {
885931

‎api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package io.sentrius.sso.websocket;
2+
3+
import org.springframework.http.server.ServerHttpRequest;
4+
import org.springframework.http.server.ServerHttpResponse;
5+
import org.springframework.web.socket.WebSocketHandler;
6+
import org.springframework.web.socket.server.HandshakeInterceptor;
7+
8+
import java.util.Map;
9+
10+
public class AgentHandshakeInterceptor implements HandshakeInterceptor {
11+
12+
@Override
13+
public boolean beforeHandshake(ServerHttpRequest request,
14+
ServerHttpResponse response,
15+
WebSocketHandler wsHandler,
16+
Map<String, Object> attributes) {
17+
18+
String path = request.getURI().getPath(); // e.g. /api/v1/agents/ws/agent-123
19+
String[] segments = path.split("/");
20+
String agentId = segments[segments.length - 1]; // assumes agentId is at the end
21+
22+
attributes.put("agentId", agentId);
23+
return true;
24+
}
25+
26+
@Override
27+
public void afterHandshake(ServerHttpRequest request,
28+
ServerHttpResponse response,
29+
WebSocketHandler wsHandler,
30+
Exception exception) {
31+
// no-op
32+
}
33+
}

‎api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package io.sentrius.sso.websocket;
2+
3+
import java.net.URI;
4+
import io.sentrius.sso.locator.KubernetesAgentLocator;
5+
import lombok.RequiredArgsConstructor;
6+
import org.springframework.stereotype.Component;
7+
import org.springframework.web.reactive.socket.WebSocketHandler;
8+
import org.springframework.web.reactive.socket.WebSocketMessage;
9+
import org.springframework.web.reactive.socket.WebSocketSession;
10+
import org.springframework.web.reactive.socket.client.ReactorNettyWebSocketClient;
11+
import reactor.core.publisher.Mono;
12+
13+
@Component
14+
@RequiredArgsConstructor
15+
public class AgentWebSocketProxyHandler implements WebSocketHandler {
16+
17+
private final KubernetesAgentLocator agentLocator;
18+
19+
@Override
20+
public Mono<Void> handle(WebSocketSession clientSession) {
21+
String agentId = (String) clientSession.getAttributes().get("agentId");
22+
URI agentUri = agentLocator.resolveWebSocketUri(agentId);
23+
24+
ReactorNettyWebSocketClient proxyClient = new ReactorNettyWebSocketClient();
25+
26+
return proxyClient.execute(agentUri, agentSession -> {
27+
// Forward messages from client to agent
28+
Mono<Void> clientToAgent = clientSession.receive()
29+
.map(WebSocketMessage::getPayload)
30+
.map(dataBuffer -> agentSession.binaryMessage(factory -> dataBuffer))
31+
.as(agentSession::send);
32+
33+
// Forward messages from agent to client
34+
Mono<Void> agentToClient = agentSession.receive()
35+
.map(WebSocketMessage::getPayload)
36+
.map(dataBuffer -> clientSession.binaryMessage(factory -> dataBuffer))
37+
.as(clientSession::send);
38+
39+
// Run both directions in parallel, complete when both are done
40+
return Mono.zip(clientToAgent, agentToClient).then();
41+
});
42+
}
43+
}

0 commit comments

Comments
 (0)