Skip to content

Improved security and injection, uses AgentCore session header #752

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ashakirin wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Improved security and injection, uses AgentCore session header #752

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions apps/java-spring-ai-agents/aiagent/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,14 +53,14 @@
<!-- AgentCore Code Interpreter -->
<dependency>
<groupId>org.springaicommunity</groupId>
<artifactId>spring-ai-bedrock-agentcore-codeinterpreter</artifactId>
<version>1.0.0-RC6</version>
<artifactId>spring-ai-agentcore-code-interpreter</artifactId>
<version>1.0.0-RC8</version>
</dependency>
<!-- AgentCore Browser -->
<dependency>
<groupId>org.springaicommunity</groupId>
<artifactId>spring-ai-bedrock-agentcore-browser</artifactId>
<version>1.0.0-RC6</version>
<artifactId>spring-ai-agentcore-browser</artifactId>
<version>1.0.0-RC8</version>
</dependency>
<!-- Bedrock Runtime SDK for web grounding -->
<dependency>
Expand All @@ -79,14 +79,14 @@
<!-- AgentCore Memory dependencies -->
<dependency>
<groupId>org.springaicommunity</groupId>
<artifactId>spring-ai-bedrock-agentcore-memory</artifactId>
<version>1.0.0-RC6</version>
<artifactId>spring-ai-agentcore-memory</artifactId>
<version>1.0.0-RC8</version>
</dependency>
<!-- AgentCore dependencies -->
<dependency>
<groupId>org.springaicommunity</groupId>
<artifactId>spring-ai-bedrock-agentcore-runtime-starter</artifactId>
<version>1.0.0-RC6</version>
<artifactId>spring-ai-agentcore-runtime-starter</artifactId>
<version>1.0.0-RC8</version>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
Expand Down
119 changes: 44 additions & 75 deletions apps/java-spring-ai-agents/aiagent/src/main/java/com/example/agent/ChatService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,43 +1,36 @@
package com.example.agent;

import java.util.ArrayList;
import java.util.List;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springaicommunity.agentcore.annotation.AgentCoreInvocation;
import org.springaicommunity.agentcore.artifacts.ArtifactStore;
import org.springaicommunity.agentcore.artifacts.GeneratedFile;
import org.springaicommunity.agentcore.artifacts.SessionConstants;
import org.springaicommunity.agentcore.browser.BrowserArtifacts;
import org.springaicommunity.agentcore.context.AgentCoreContext;
import org.springaicommunity.agentcore.context.AgentCoreHeaders;
import org.springaicommunity.agentcore.browser.BrowserScreenshot;
import org.springaicommunity.agentcore.browser.BrowserScreenshotStore;
import org.springaicommunity.agentcore.browser.BrowserTools;
import org.springaicommunity.agentcore.codeinterpreter.CodeInterpreterFileStore;
import org.springaicommunity.agentcore.codeinterpreter.CodeInterpreterTools;
import org.springaicommunity.agentcore.codeinterpreter.GeneratedFile;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springaicommunity.agentcore.memory.longterm.AgentCoreMemory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import tools.jackson.databind.JsonNode;
import tools.jackson.databind.json.JsonMapper;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.http.MediaType;
import org.springframework.http.MediaTypeFactory;
import org.springframework.stereotype.Service;
import org.springframework.util.MimeType;
import org.springframework.util.MimeTypeUtils;
import reactor.core.publisher.Flux;
import tools.jackson.databind.json.JsonMapper;

import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

record ChatRequest(String prompt, String fileBase64, String fileName) {
public boolean hasFile() {
Expand All @@ -57,43 +50,32 @@ public class ChatService {

private final JsonMapper jsonMapper = JsonMapper.builder().build();

private final BrowserScreenshotStore screenshotStore;
private final CodeInterpreterFileStore fileStore;
private final ArtifactStore<GeneratedFile> browserArtifactStore;
private final ArtifactStore<GeneratedFile> codeInterpreterArtifactStore;

private static final String SYSTEM_PROMPT = """
You are a helpful AI agent for travel and expense management.
Be friendly, helpful, and concise in your responses.
""";

public ChatService(@Autowired(required = false) ChatMemoryRepository memoryRepository,
@Autowired(required = false) List<Advisor> ltmAdvisors,
@Autowired(required = false) VectorStore kbVectorStore,
@Autowired(required = false) WebGroundingTools webGroundingTools,
ContextAdvisor contextAdvisor,
@Autowired(required = false) @Qualifier("browserToolCallbackProvider") ToolCallbackProvider browserTools,
@Autowired(required = false) BrowserScreenshotStore browserScreenshotStore,
@Autowired(required = false) @Qualifier("codeInterpreterToolCallbackProvider") ToolCallbackProvider codeInterpreterTools,
@Autowired(required = false) CodeInterpreterFileStore codeInterpreterFileStore,
@Autowired(required = false) @Qualifier("mcpToolCallbacks") ToolCallbackProvider mcpTools,
ChatModel chatModel,
@Value("${app.ai.document.model:global.anthropic.claude-opus-4-5-20251101-v1:0}") String documentModel,
ChatClient.Builder chatClientBuilder) {
public ChatService(AgentCoreMemory agentCoreMemory,
VectorStore kbVectorStore,
WebGroundingTools webGroundingTools,
ContextAdvisor contextAdvisor,
@Qualifier("browserToolCallbackProvider") ToolCallbackProvider browserTools,
@Qualifier("codeInterpreterToolCallbackProvider") ToolCallbackProvider codeInterpreterTools,
@Qualifier("browserArtifactStore") ArtifactStore<GeneratedFile> browserArtifactStore,
@Qualifier("codeInterpreterArtifactStore") ArtifactStore<GeneratedFile> codeInterpreterArtifactStore,
@Qualifier("mcpToolCallbacks") ToolCallbackProvider mcpTools,
ChatModel chatModel,
@Value("${app.ai.document.model:global.anthropic.claude-opus-4-5-20251101-v1:0}") String documentModel,
ChatClient.Builder chatClientBuilder) {

List<Advisor> advisors = new ArrayList<>();

// Short-Term Memory (STM)
if (memoryRepository != null) {
ChatMemory chatMemory = MessageWindowChatMemory.builder()
.chatMemoryRepository(memoryRepository)
.build();
advisors.add(MessageChatMemoryAdvisor.builder(chatMemory).build());
logger.info("Memory enabled");
}

// Long-Term Memory (LTM)
if (ltmAdvisors != null && !ltmAdvisors.isEmpty()) {
advisors.addAll(ltmAdvisors);
logger.info("Advisors enabled: {} advisors", ltmAdvisors.size());
if (advisors.size() > 0) {
advisors.addAll(agentCoreMemory.advisors);
logger.info("Advisors enabled: {} advisors", agentCoreMemory.advisors);
}

// Knowledge Base (RAG)
Expand All @@ -113,19 +95,15 @@ public ChatService(@Autowired(required = false) ChatMemoryRepository memoryRepos
logger.info("Web Grounding enabled");
}

// Browser
this.screenshotStore = browserScreenshotStore;

// Tool Callback Providers
this.browserArtifactStore = browserArtifactStore;
List<ToolCallbackProvider> toolCallbackProviders = new ArrayList<>();
if (browserTools != null) {
toolCallbackProviders.add(browserTools);
logger.info("Browser enabled");
}

// Code Interpreter
this.fileStore = codeInterpreterFileStore;

this.codeInterpreterArtifactStore = codeInterpreterArtifactStore;
if (codeInterpreterTools != null) {
toolCallbackProviders.add(codeInterpreterTools);
logger.info("Code Interpreter enabled");
Expand Down Expand Up @@ -169,38 +147,29 @@ private Flux<String> chat(String prompt, String sessionId) {
.stream().content()
.concatWith(Flux.defer(() -> appendGeneratedFiles(sessionId)))
.concatWith(Flux.defer(() -> appendScreenshots(sessionId)))
.contextWrite(ctx -> ctx.put(CodeInterpreterTools.SESSION_ID_CONTEXT_KEY, sessionId))
.contextWrite(ctx -> ctx.put(BrowserTools.SESSION_ID_CONTEXT_KEY, sessionId));
.contextWrite(ctx -> ctx.put(SessionConstants.SESSION_ID_KEY, sessionId));
}

private String getSessionId(AgentCoreContext context) {
String authHeader = context.getHeader(AgentCoreHeaders.AUTHORIZATION);
if (authHeader != null && authHeader.startsWith("Bearer ")) {
String jwt = authHeader.replace("Bearer ", "");
String payload = new String(Base64.getUrlDecoder().decode(jwt.split("\\.")[1]));
JsonNode claims = jsonMapper.readTree(payload);
String sub = claims.get("sub").asString().replace("-", "").substring(0, 25);
return sub + ":" + claims.get("auth_time").asString();
}
return (authHeader != null && !authHeader.isBlank()) ? authHeader : "default-user";
return ConversationIdResolver.resolve(context);
}

private Flux<String> appendScreenshots(String sessionId) {
if (screenshotStore == null) {
if (browserArtifactStore == null) {
return Flux.empty();
}
List<BrowserScreenshot> screenshots = screenshotStore.retrieve(sessionId);
List<GeneratedFile> screenshots = browserArtifactStore.retrieve(sessionId);
if (screenshots == null || screenshots.isEmpty()) {
return Flux.empty();
}
return Flux.just(formatScreenshotsAsMarkdown(screenshots));
}

private String formatScreenshotsAsMarkdown(List<BrowserScreenshot> screenshots) {
private String formatScreenshotsAsMarkdown(List<GeneratedFile> screenshots) {
StringBuilder sb = new StringBuilder();
for (BrowserScreenshot screenshot : screenshots) {
for (GeneratedFile screenshot : screenshots) {
sb.append("\n\n![Screenshot of ")
.append(screenshot.url())
.append(BrowserArtifacts.url(screenshot))
.append("](")
.append(screenshot.toDataUrl())
.append(")");
Expand All @@ -209,10 +178,10 @@ private String formatScreenshotsAsMarkdown(List<BrowserScreenshot> screenshots)
}

private Flux<String> appendGeneratedFiles(String sessionId) {
if (fileStore == null) {
if (codeInterpreterArtifactStore == null) {
return Flux.empty();
}
List<GeneratedFile> files = fileStore.retrieve(sessionId);
List<GeneratedFile> files = codeInterpreterArtifactStore.retrieve(sessionId);
if (files == null || files.isEmpty()) {
return Flux.empty();
}
Expand Down
45 changes: 45 additions & 0 deletions ...java-spring-ai-agents/aiagent/src/main/java/com/example/agent/ConversationIdResolver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
package com.example.agent;

import org.springaicommunity.agentcore.context.AgentCoreContext;
import org.springaicommunity.agentcore.context.AgentCoreHeaders;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tools.jackson.databind.JsonNode;
import tools.jackson.databind.json.JsonMapper;

import java.util.Base64;
import java.util.UUID;

/**
* Utility for extracting conversation ID from AgentCore context.
* Format: userId:sessionId (authenticated) or sessionId (anonymous)
*/
public final class ConversationIdResolver {

private static final Logger logger = LoggerFactory.getLogger(ConversationIdResolver.class);
private static final JsonMapper jsonMapper = JsonMapper.builder().build();

private ConversationIdResolver() {}

public static String resolve(AgentCoreContext context) {
String sessionId = context.getHeader(AgentCoreHeaders.SESSION_ID);
if (sessionId == null || sessionId.isBlank()) {
sessionId = UUID.randomUUID().toString();
}

String authHeader = context.getHeader(AgentCoreHeaders.AUTHORIZATION);
if (authHeader != null && authHeader.startsWith("Bearer ")) {
try {
String jwt = authHeader.substring(7);
String payload = new String(Base64.getUrlDecoder().decode(jwt.split("\\.")[1]));
JsonNode claims = jsonMapper.readTree(payload);
String userId = claims.get("sub").asString();
return userId + ":" + sessionId;
} catch (Exception e) {
logger.debug("JWT parsing failed, using sessionId only", e);
}
}

return sessionId;
}
}
5 changes: 5 additions & 0 deletions apps/java-spring-ai-agents/aiagent/src/main/resources/application.properties
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,11 @@ spring.ai.bedrock.converse.chat.timeout=120s
spring.ai.bedrock.converse.chat.options.max-tokens=4096
spring.ai.bedrock.converse.chat.options.model=global.anthropic.claude-sonnet-4-5-20250929-v1:0
spring.ai.bedrock.converse.chat.options.temperature=0.7
# Knowledge Base
spring.ai.vectorstore.bedrock-knowledge-base.knowledge-base-id=${BEDROCK_KNOWLEDGE_BASE_ID:your-knowledge-base-id-here}
# AgentCore memory
agentcore.memory.memory-id=${AGENTCORE_MEMORY_ID:your-memory-id-here}
agentcore.memory.long-term.auto-discovery=true
# AgentCore Browser - tool descriptions
agentcore.browser.browse-url-description=Browse a web page and extract its text content. Returns the page title and body text. Use this to read and extract data from websites. For interactive sites, combine with fillForm and clickElement to navigate, then call browseUrl again to read the results.
agentcore.browser.screenshot-description=Take a screenshot of a web page for the user to see. Does NOT return page content to you. Use browseUrl to extract data first, then takeScreenshot for visual evidence.
Expand Down
39 changes: 18 additions & 21 deletions apps/java-spring-ai-agents/aiagent/src/main/resources/static/auth.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,23 @@
// Authentication management
const AUTH_KEY = 'agentcore_auth';
const SESSION_KEY = 'agentcore_session_id';

let currentAuth = null;

// Session ID management
function getSessionId() {
let sessionId = sessionStorage.getItem(SESSION_KEY);
if (!sessionId) {
sessionId = crypto.randomUUID();
sessionStorage.setItem(SESSION_KEY, sessionId);
}
return sessionId;
}

function clearSessionId() {
sessionStorage.removeItem(SESSION_KEY);
}

function saveAuth(auth) {
localStorage.setItem(AUTH_KEY, JSON.stringify(auth));
currentAuth = auth;
Expand All @@ -19,13 +34,13 @@ function loadAuth() {

function clearAuth() {
localStorage.removeItem(AUTH_KEY);
clearSessionId();
currentAuth = null;
}

function isAuthenticated() {
const auth = loadAuth();
if (!auth || !auth.expiresAt) return false;
if (auth.authType !== 'simple' && !auth.accessToken) return false;
if (!auth || !auth.expiresAt || !auth.accessToken) return false;
return !isSessionExpired(auth);
}

Expand All @@ -34,25 +49,7 @@ function isSessionExpired(auth) {
}

async function authenticateUser(username, password, config) {
if (config.authType === 'cognito') {
return authenticateCognito(username, password, config);
}
return authenticateSimple(username, password);
}

async function authenticateSimple(username, password) {
if (!username || username.trim() === '') {
throw new Error('Username is required');
}

const auth = {
username: username,
expiresAt: Date.now() + (24 * 60 * 60 * 1000),
authType: 'simple'
};

saveAuth(auth);
return auth;
return authenticateCognito(username, password, config);
}

async function authenticateCognito(username, password, config) {
Expand Down
3 changes: 2 additions & 1 deletion apps/java-spring-ai-agents/aiagent/src/main/resources/static/chat.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,7 +317,8 @@ async function sendMessage(message, config, auth) {

const headers = {
'Content-Type': 'application/json',
'Accept': 'text/plain, text/event-stream'
'Accept': 'text/plain, text/event-stream',
'X-Amzn-Bedrock-AgentCore-Runtime-Session-Id': getSessionId()
};

if (config.authType === 'cognito' || config.mode === 'aws') {
Expand Down
Loading