diff --git a/agentscope-examples/agui/src/main/java/io/agentscope/examples/agui/config/AgentConfiguration.java b/agentscope-examples/agui/src/main/java/io/agentscope/examples/agui/config/AgentConfiguration.java
index dae080f07..21204504d 100644
--- a/agentscope-examples/agui/src/main/java/io/agentscope/examples/agui/config/AgentConfiguration.java
+++ b/agentscope-examples/agui/src/main/java/io/agentscope/examples/agui/config/AgentConfiguration.java
@@ -90,8 +90,10 @@ private Agent createDefaultAgent() {
+ "You can help users with various tasks including weather queries "
+ "and calculations. Be concise and helpful in your responses.")
.model(
- DashScopeChatModel.builder().apiKey(apiKey).modelName("qwen-plus").stream(
- true)
+ DashScopeChatModel.builder()
+ .apiKey(apiKey)
+ .modelName("qwen3.6-plus")
+ .stream(true)
.enableThinking(false)
.formatter(new DashScopeChatFormatter())
.build())
diff --git a/agentscope-examples/agui/src/main/java/io/agentscope/examples/agui/config/WebFluxCodecConfiguration.java b/agentscope-examples/agui/src/main/java/io/agentscope/examples/agui/config/WebFluxCodecConfiguration.java
new file mode 100644
index 000000000..f1d2c1b38
--- /dev/null
+++ b/agentscope-examples/agui/src/main/java/io/agentscope/examples/agui/config/WebFluxCodecConfiguration.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.examples.agui.config;
+
+import org.springframework.boot.http.codec.CodecCustomizer;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+/**
+ * Overrides the default WebFlux codec buffer limit (256KB) so that AG-UI requests
+ * carrying long conversation history or large message payloads are not rejected
+ * with {@code DataBufferLimitException: Exceeded limit on max bytes to buffer : 262144}.
+ *
+ *
async function sendMessage() {
const text = input.value.trim();
- if (!text || isRunning) return;
+ // Require either text or at least one attachment to send.
+ if ((!text && pendingAttachments.length === 0) || isRunning) return;
+
+ // Snapshot the attachments for this message, then clear the pending list.
+ const attachmentsForMessage = pendingAttachments;
+ pendingAttachments = [];
+ renderAttachmentsPreview();
input.value = '';
isRunning = true;
@@ -391,9 +602,33 @@
AgentScope AG-UI Demo
stopBtn.style.display = 'inline-block';
setStatus('running', 'Running...');
- // Add user message
- appendMessage('user', text);
- messageHistory.push({ id: 'msg-' + Date.now(), role: 'user', content: text });
+ // Build the AG-UI UserMessage content.
+ // - Plain text only => content is a string (backward compatible)
+ // - Any attachment present => content is an array of parts
+ let userContent;
+ if (attachmentsForMessage.length === 0) {
+ userContent = text;
+ } else {
+ userContent = [];
+ if (text) {
+ userContent.push({ type: 'text', text: text });
+ }
+ for (const att of attachmentsForMessage) {
+ userContent.push({
+ type: 'image',
+ source: {
+ type: att.source.type,
+ value: att.source.value,
+ mimeType: att.source.mimeType
+ }
+ });
+ }
+ }
+
+ // Render user message in UI (with image thumbnails if any).
+ const previewUrls = attachmentsForMessage.map(a => a.previewUrl);
+ appendMessage('user', text || '(image)', false, previewUrls.length ? previewUrls : null);
+ messageHistory.push({ id: 'msg-' + Date.now(), role: 'user', content: userContent });
// Show typing indicator
showTypingIndicator();
@@ -534,6 +769,37 @@
AgentScope AG-UI Demo
sendBtn.addEventListener('click', sendMessage);
stopBtn.addEventListener('click', stopGeneration);
+ // Attachment: upload local image files
+ uploadBtn.addEventListener('click', () => fileInput.click());
+ fileInput.addEventListener('change', async (e) => {
+ const files = Array.from(e.target.files || []);
+ if (files.length > 0) {
+ await handleFileSelection(files);
+ }
+ // Reset so selecting the same file again still triggers change.
+ fileInput.value = '';
+ });
+
+ // Attachment: add image URL
+ urlBtn.addEventListener('click', handleAddImageUrl);
+
+ // Attachment: paste image from clipboard directly into the input
+ input.addEventListener('paste', async (e) => {
+ const items = e.clipboardData?.items;
+ if (!items) return;
+ const imageFiles = [];
+ for (const item of items) {
+ if (item.kind === 'file' && item.type.startsWith('image/')) {
+ const file = item.getAsFile();
+ if (file) imageFiles.push(file);
+ }
+ }
+ if (imageFiles.length > 0) {
+ e.preventDefault();
+ await handleFileSelection(imageFiles);
+ }
+ });
+
// Focus input on load
input.focus();
diff --git a/agentscope-extensions/agentscope-extensions-agui/pom.xml b/agentscope-extensions/agentscope-extensions-agui/pom.xml
index 5b896643c..8b68a9bc0 100644
--- a/agentscope-extensions/agentscope-extensions-agui/pom.xml
+++ b/agentscope-extensions/agentscope-extensions-agui/pom.xml
@@ -37,6 +37,34 @@
providedtrue
+
+
+
+ com.fasterxml.jackson.core
+ jackson-databind
+ provided
+ true
+
+
+
+
+ io.projectreactor
+ reactor-core
+ provided
+ true
+
+
+
+
+ org.junit.jupiter
+ junit-jupiter
+ test
+
+
+ org.mockito
+ mockito-core
+ test
+
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/AguiRequestContext.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/AguiRequestContext.java
new file mode 100644
index 000000000..ffc52dfdd
--- /dev/null
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/AguiRequestContext.java
@@ -0,0 +1,173 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.core.agui;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+
+/**
+ * Thread-local context for passing HTTP request metadata to agent factories.
+ *
+ *
This class provides a mechanism for agent factory methods (e.g., {@code @AgentScopeBean}
+ * annotated methods) to access HTTP request headers and query parameters from the current
+ * AGUI request. The context is set before agent resolution and cleared after processing.
+ *
+ *
Thread safety: The context is bound to the current thread via {@link ThreadLocal}.
+ * It is set before {@code processor.process()} and cleared in a {@code finally} block,
+ * ensuring no leaks across requests in thread pools.
+ *
+ *
Null safety: {@link #current()} never returns {@code null}. When no context
+ * is set (e.g., in tests or non-AGUI paths), it returns an empty context where all
+ * getter methods return {@code null} or empty collections.
+ */
+public class AguiRequestContext {
+
+ private static final ThreadLocal HOLDER = new ThreadLocal<>();
+
+ private static final AguiRequestContext EMPTY =
+ new AguiRequestContext(Collections.emptyMap(), Collections.emptyMap());
+
+ private final Map> headers;
+ private final Map> params;
+
+ private AguiRequestContext(
+ Map> headers, Map> params) {
+ // Defensive copy with case-insensitive keys for headers
+ TreeMap> headersCopy = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
+ headersCopy.putAll(headers);
+ this.headers = Collections.unmodifiableMap(headersCopy);
+ this.params = Collections.unmodifiableMap(params);
+ }
+
+ /**
+ * Initialize the request context for the current thread.
+ *
+ *
Must be paired with {@link #clear()} in a {@code finally} block.
+ *
+ * @param headers HTTP request headers
+ * @param params HTTP query parameters
+ */
+ public static void init(Map> headers, Map> params) {
+ HOLDER.set(
+ new AguiRequestContext(
+ headers != null ? headers : Collections.emptyMap(),
+ params != null ? params : Collections.emptyMap()));
+ }
+
+ /**
+ * Get the current request context.
+ *
+ *
Never returns {@code null}. Returns an empty context when no context is set.
+ *
+ * @return The current request context, or an empty context
+ */
+ public static AguiRequestContext current() {
+ AguiRequestContext ctx = HOLDER.get();
+ return ctx != null ? ctx : EMPTY;
+ }
+
+ /**
+ * Clear the request context for the current thread.
+ *
+ *
Must be called in a {@code finally} block after {@link #init}.
+ */
+ public static void clear() {
+ HOLDER.remove();
+ }
+
+ // --- Header methods ---
+
+ /**
+ * Get the first value of the specified HTTP header.
+ *
+ *
Header name lookup is case-insensitive per HTTP specification.
+ *
+ * @param name The header name
+ * @return The first header value, or {@code null} if not present
+ */
+ public String getHeader(String name) {
+ List values = headers.get(name);
+ return (values != null && !values.isEmpty()) ? values.get(0) : null;
+ }
+
+ /**
+ * Get all values of the specified HTTP header.
+ *
+ * @param name The header name
+ * @return All header values, or an empty list if not present
+ */
+ public List getHeaders(String name) {
+ List values = headers.get(name);
+ return values != null ? values : Collections.emptyList();
+ }
+
+ /**
+ * Get all HTTP headers.
+ *
+ * @return An unmodifiable map of all headers
+ */
+ public Map> getAllHeaders() {
+ return headers;
+ }
+
+ // --- Parameter methods ---
+
+ /**
+ * Get the first value of the specified query parameter.
+ *
+ * @param name The parameter name
+ * @return The first parameter value, or {@code null} if not present
+ */
+ public String getParameter(String name) {
+ List values = params.get(name);
+ return (values != null && !values.isEmpty()) ? values.get(0) : null;
+ }
+
+ /**
+ * Get all values of the specified query parameter.
+ *
+ * @param name The parameter name
+ * @return All parameter values, or an empty list if not present
+ */
+ public List getParameters(String name) {
+ List values = params.get(name);
+ return values != null ? values : Collections.emptyList();
+ }
+
+ /**
+ * Get all query parameters.
+ *
+ * @return An unmodifiable map of all parameters
+ */
+ public Map> getAllParameters() {
+ return params;
+ }
+}
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/converter/AguiMessageConverter.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/converter/AguiMessageConverter.java
index 78ee08a22..f5ebb273b 100644
--- a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/converter/AguiMessageConverter.java
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/converter/AguiMessageConverter.java
@@ -16,15 +16,25 @@
package io.agentscope.core.agui.converter;
import com.fasterxml.jackson.core.type.TypeReference;
+import io.agentscope.core.agui.model.AguiContentPart;
+import io.agentscope.core.agui.model.AguiContentSource;
+import io.agentscope.core.agui.model.AguiDataSource;
import io.agentscope.core.agui.model.AguiFunctionCall;
+import io.agentscope.core.agui.model.AguiImageContent;
import io.agentscope.core.agui.model.AguiMessage;
+import io.agentscope.core.agui.model.AguiTextContent;
import io.agentscope.core.agui.model.AguiToolCall;
+import io.agentscope.core.agui.model.AguiUrlSource;
+import io.agentscope.core.message.Base64Source;
import io.agentscope.core.message.ContentBlock;
+import io.agentscope.core.message.ImageBlock;
import io.agentscope.core.message.Msg;
import io.agentscope.core.message.MsgRole;
+import io.agentscope.core.message.Source;
import io.agentscope.core.message.TextBlock;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.message.ToolUseBlock;
+import io.agentscope.core.message.URLSource;
import io.agentscope.core.util.JsonException;
import io.agentscope.core.util.JsonUtils;
import java.util.ArrayList;
@@ -54,17 +64,40 @@ public Msg toMsg(AguiMessage aguiMessage) {
MsgRole role = convertRole(aguiMessage.getRole());
List blocks = new ArrayList<>();
- // Add text content if present
- if (aguiMessage.getContent() != null && !aguiMessage.getContent().isEmpty()) {
- if (aguiMessage.isToolMessage() && aguiMessage.getToolCallId() != null) {
- // For tool messages, wrap content in ToolResultBlock
- blocks.add(
- ToolResultBlock.of(
- aguiMessage.getToolCallId(),
- null,
- TextBlock.builder().text(aguiMessage.getContent()).build()));
- } else {
- blocks.add(TextBlock.builder().text(aguiMessage.getContent()).build());
+ if (aguiMessage.isMultimodal()) {
+ // Handle multimodal content parts (text, image, etc.)
+ for (AguiContentPart part : aguiMessage.getContentParts()) {
+ if (part instanceof AguiTextContent textContent) {
+ if (aguiMessage.isToolMessage() && aguiMessage.getToolCallId() != null) {
+ blocks.add(
+ ToolResultBlock.of(
+ aguiMessage.getToolCallId(),
+ null,
+ TextBlock.builder().text(textContent.text()).build()));
+ } else {
+ blocks.add(TextBlock.builder().text(textContent.text()).build());
+ }
+ } else if (part instanceof AguiImageContent imageContent) {
+ blocks.add(
+ ImageBlock.builder()
+ .source(convertSource(imageContent.source()))
+ .build());
+ }
+ // Other content types (audio, video, document) are silently skipped
+ }
+ } else {
+ // Add text content if present (plain string format)
+ if (aguiMessage.getContent() != null && !aguiMessage.getContent().isEmpty()) {
+ if (aguiMessage.isToolMessage() && aguiMessage.getToolCallId() != null) {
+ // For tool messages, wrap content in ToolResultBlock
+ blocks.add(
+ ToolResultBlock.of(
+ aguiMessage.getToolCallId(),
+ null,
+ TextBlock.builder().text(aguiMessage.getContent()).build()));
+ } else {
+ blocks.add(TextBlock.builder().text(aguiMessage.getContent()).build());
+ }
}
}
@@ -232,4 +265,22 @@ private String serializeArguments(Map arguments) {
return "{}";
}
}
+
+ /**
+ * Convert an AG-UI content source to an AgentScope Source.
+ *
+ * @param source The AG-UI content source
+ * @return The converted AgentScope Source
+ */
+ private Source convertSource(AguiContentSource source) {
+ if (source instanceof AguiDataSource dataSource) {
+ return Base64Source.builder()
+ .data(dataSource.value())
+ .mediaType(dataSource.mimeType())
+ .build();
+ } else if (source instanceof AguiUrlSource urlSource) {
+ return URLSource.builder().url(urlSource.value()).build();
+ }
+ throw new IllegalArgumentException("Unknown content source type: " + source.getClass());
+ }
}
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiContentPart.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiContentPart.java
new file mode 100644
index 000000000..bd6b387bf
--- /dev/null
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiContentPart.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.core.agui.model;
+
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+
+/**
+ * Base interface for multimodal content parts in the AG-UI protocol.
+ *
+ *
The AG-UI protocol supports multimodal messages where the {@code content} field
+ * can be an array of typed content parts. Each part specifies a {@code type} discriminator
+ * that determines the concrete content type.
+ *
+ * @see AguiTextContent
+ * @see AguiImageContent
+ */
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
+@JsonSubTypes({
+ @JsonSubTypes.Type(value = AguiTextContent.class, name = "text"),
+ @JsonSubTypes.Type(value = AguiImageContent.class, name = "image")
+})
+public sealed interface AguiContentPart permits AguiTextContent, AguiImageContent {}
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiContentSource.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiContentSource.java
new file mode 100644
index 000000000..bcc7b8cf8
--- /dev/null
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiContentSource.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.core.agui.model;
+
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+
+/**
+ * Base interface for content source types in the AG-UI protocol.
+ *
+ *
Sources can be either inline data (base64-encoded) or URL references.
+ *
+ * @see AguiDataSource
+ * @see AguiUrlSource
+ */
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
+@JsonSubTypes({
+ @JsonSubTypes.Type(value = AguiDataSource.class, name = "data"),
+ @JsonSubTypes.Type(value = AguiUrlSource.class, name = "url")
+})
+public sealed interface AguiContentSource permits AguiDataSource, AguiUrlSource {}
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiDataSource.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiDataSource.java
new file mode 100644
index 000000000..a426d5bba
--- /dev/null
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiDataSource.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.core.agui.model;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.Objects;
+
+/**
+ * Represents an inline base64-encoded data source in the AG-UI protocol.
+ *
+ *
JSON format: {@code {"type": "data", "value": "", "mimeType": "image/png"}}
+ */
+public record AguiDataSource(String value, String mimeType) implements AguiContentSource {
+
+ @JsonCreator
+ public AguiDataSource(
+ @JsonProperty("value") String value, @JsonProperty("mimeType") String mimeType) {
+ this.value = Objects.requireNonNull(value, "value cannot be null");
+ this.mimeType = Objects.requireNonNull(mimeType, "mimeType cannot be null");
+ }
+}
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiImageContent.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiImageContent.java
new file mode 100644
index 000000000..039667b60
--- /dev/null
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiImageContent.java
@@ -0,0 +1,59 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.core.agui.model;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Represents an image content part in the AG-UI protocol's multimodal message format.
+ *
+ *
Messages are the primary communication unit in the AG-UI protocol.
* They contain content, role information, and optionally tool calls or tool call IDs.
*
+ *
The {@code content} field supports two formats:
+ *
+ *
Plain text string - for simple text messages (backward compatible)
+ *
Array of content parts - for multimodal messages containing text, images, etc.
+ *
+ *
*
Message roles:
*
*
user - Messages from the user
@@ -35,33 +45,40 @@
*
tool - Tool execution results
*
*/
+@JsonDeserialize(using = AguiMessageDeserializer.class)
public class AguiMessage {
private final String id;
private final String role;
private final String content;
+ private final List contentParts;
private final List toolCalls;
private final String toolCallId;
/**
- * Creates a new AguiMessage.
+ * Creates a new AguiMessage with full control over all fields.
+ *
+ *
At most one of {@code content} and {@code contentParts} should be non-null.
*
* @param id The unique message ID
* @param role The message role (user, assistant, system, tool)
- * @param content The message content
+ * @param content The plain text content (for simple messages)
+ * @param contentParts The multimodal content parts (for multimodal messages)
* @param toolCalls Tool calls for assistant messages (optional)
* @param toolCallId Tool call ID for tool messages (optional)
*/
- @JsonCreator
public AguiMessage(
- @JsonProperty("id") String id,
- @JsonProperty("role") String role,
- @JsonProperty("content") String content,
- @JsonProperty("toolCalls") List toolCalls,
- @JsonProperty("toolCallId") String toolCallId) {
+ String id,
+ String role,
+ String content,
+ List contentParts,
+ List toolCalls,
+ String toolCallId) {
this.id = Objects.requireNonNull(id, "id cannot be null");
this.role = Objects.requireNonNull(role, "role cannot be null");
this.content = content;
+ this.contentParts =
+ contentParts != null ? Collections.unmodifiableList(contentParts) : null;
this.toolCalls =
toolCalls != null
? Collections.unmodifiableList(toolCalls)
@@ -69,6 +86,121 @@ public AguiMessage(
this.toolCallId = toolCallId;
}
+ /**
+ * Jackson-friendly creator constructor used by both Jackson 2.x and Jackson 3.x.
+ *
+ *
The core Jackson annotations ({@link JsonCreator}, {@link JsonProperty}) share
+ * the same package ({@code com.fasterxml.jackson.annotation}) across Jackson 2 and
+ * Jackson 3, so this constructor is discovered by both versions and serves as a
+ * portable entry point when the Jackson 2-only {@link JsonDeserialize} on this class
+ * is ignored by Jackson 3 (e.g. Spring Framework 7 / Spring Boot 4 use Jackson 3
+ * under the {@code tools.jackson.*} package).
+ *
+ *
The {@code content} field is a union type in the AG-UI protocol
+ * ({@code string | InputContent[]}). This constructor accepts it as {@link Object}
+ * and dispatches by runtime type:
+ *
+ *
{@link String} - stored as plain text content
+ *
{@link List} - converted to {@link AguiContentPart} entries by inspecting
+ * the {@code type} discriminator on each element map
+ *
+ */
+ @JsonCreator
+ @SuppressWarnings("unchecked")
+ static AguiMessage fromJson(
+ @JsonProperty("id") String id,
+ @JsonProperty("role") String role,
+ @JsonProperty("content") Object content,
+ @JsonProperty("toolCalls") List toolCalls,
+ @JsonProperty("toolCallId") String toolCallId) {
+ String textContent = null;
+ List parts = null;
+
+ if (content instanceof String text) {
+ textContent = text;
+ } else if (content instanceof List list) {
+ parts = new ArrayList<>(list.size());
+ for (Object element : list) {
+ if (element instanceof AguiContentPart part) {
+ parts.add(part);
+ } else if (element instanceof Map map) {
+ parts.add(toContentPart((Map) map));
+ } else if (element != null) {
+ throw new IllegalArgumentException(
+ "Unsupported content part element type: "
+ + element.getClass().getName());
+ }
+ }
+ } else if (content != null) {
+ throw new IllegalArgumentException(
+ "Unsupported content type: " + content.getClass().getName());
+ }
+
+ return new AguiMessage(id, role, textContent, parts, toolCalls, toolCallId);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static AguiContentPart toContentPart(Map map) {
+ Object type = map.get("type");
+ if (type == null) {
+ throw new IllegalArgumentException(
+ "Content part is missing required 'type' discriminator: " + map);
+ }
+ String typeStr = type.toString();
+ return switch (typeStr) {
+ case "text" -> new AguiTextContent(asString(map.get("text")));
+ case "image" ->
+ new AguiImageContent(
+ toContentSource((Map) map.get("source")),
+ (Map) map.get("metadata"));
+ default ->
+ throw new IllegalArgumentException("Unsupported content part type: " + typeStr);
+ };
+ }
+
+ private static AguiContentSource toContentSource(Map map) {
+ if (map == null) {
+ throw new IllegalArgumentException("Image content is missing required 'source' field");
+ }
+ Object type = map.get("type");
+ if (type == null) {
+ throw new IllegalArgumentException(
+ "Content source is missing required 'type' discriminator: " + map);
+ }
+ String typeStr = type.toString();
+ return switch (typeStr) {
+ case "url" ->
+ new AguiUrlSource(asString(map.get("value")), asString(map.get("mimeType")));
+ case "data" ->
+ new AguiDataSource(asString(map.get("value")), asString(map.get("mimeType")));
+ default ->
+ throw new IllegalArgumentException(
+ "Unsupported content source type: " + typeStr);
+ };
+ }
+
+ private static String asString(Object value) {
+ return value == null ? null : value.toString();
+ }
+
+ /**
+ * Creates a new AguiMessage with plain text content (backward compatible constructor).
+ *
+ * @param id The unique message ID
+ * @param role The message role (user, assistant, system, tool)
+ * @param content The message content
+ * @param toolCalls Tool calls for assistant messages (optional)
+ * @param toolCallId Tool call ID for tool messages (optional)
+ */
+ public AguiMessage(
+ String id,
+ String role,
+ String content,
+ List toolCalls,
+ String toolCallId) {
+ this(id, role, content, null, toolCalls, toolCallId);
+ }
+
/**
* Creates a simple user message.
*
@@ -114,11 +246,25 @@ public static AguiMessage toolMessage(String id, String toolCallId, String conte
return new AguiMessage(id, "tool", content, null, toolCallId);
}
+ /**
+ * Creates a multimodal message with content parts.
+ *
+ * @param id The message ID
+ * @param role The message role
+ * @param contentParts The multimodal content parts
+ * @return A new multimodal message
+ */
+ public static AguiMessage multimodalMessage(
+ String id, String role, List contentParts) {
+ return new AguiMessage(id, role, null, contentParts, null, null);
+ }
+
/**
* Get the message ID.
*
* @return The message ID
*/
+ @JsonProperty("id")
public String getId() {
return id;
}
@@ -128,24 +274,47 @@ public String getId() {
*
* @return The role (user, assistant, system, tool)
*/
+ @JsonProperty("role")
public String getRole() {
return role;
}
/**
- * Get the message content.
+ * Get the plain text message content.
*
- * @return The content, may be null
+ * @return The content string, or null if this is a multimodal message
*/
+ @JsonProperty("content")
public String getContent() {
return content;
}
+ /**
+ * Get the multimodal content parts.
+ *
+ * @return The content parts list, or null if this is a plain text message
+ */
+ @JsonIgnore
+ public List getContentParts() {
+ return contentParts;
+ }
+
+ /**
+ * Check if this message contains multimodal content parts.
+ *
+ * @return true if content parts are present
+ */
+ @JsonIgnore
+ public boolean isMultimodal() {
+ return contentParts != null && !contentParts.isEmpty();
+ }
+
/**
* Get the tool calls (for assistant messages).
*
* @return The tool calls as an immutable list, empty if none
*/
+ @JsonProperty("toolCalls")
public List getToolCalls() {
return toolCalls;
}
@@ -155,6 +324,7 @@ public List getToolCalls() {
*
* @return The tool call ID, or null if not a tool message
*/
+ @JsonProperty("toolCallId")
public String getToolCallId() {
return toolCallId;
}
@@ -212,7 +382,9 @@ public String toString() {
+ role
+ "', content='"
+ content
- + "', toolCalls="
+ + "', contentParts="
+ + contentParts
+ + ", toolCalls="
+ toolCalls
+ ", toolCallId='"
+ toolCallId
@@ -227,12 +399,13 @@ public boolean equals(Object o) {
return Objects.equals(id, that.id)
&& Objects.equals(role, that.role)
&& Objects.equals(content, that.content)
+ && Objects.equals(contentParts, that.contentParts)
&& Objects.equals(toolCalls, that.toolCalls)
&& Objects.equals(toolCallId, that.toolCallId);
}
@Override
public int hashCode() {
- return Objects.hash(id, role, content, toolCalls, toolCallId);
+ return Objects.hash(id, role, content, contentParts, toolCalls, toolCallId);
}
}
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiMessageDeserializer.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiMessageDeserializer.java
new file mode 100644
index 000000000..a0f8f35d9
--- /dev/null
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiMessageDeserializer.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.core.agui.model;
+
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JsonDeserializer;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Custom Jackson deserializer for {@link AguiMessage} that handles the union type
+ * {@code content: string | InputContent[]}.
+ *
+ *
The AG-UI protocol allows the {@code content} field to be either a plain text string
+ * or an array of multimodal content parts. This deserializer inspects the JSON token type
+ * and routes accordingly.
+ */
+public class AguiMessageDeserializer extends JsonDeserializer {
+
+ @Override
+ public AguiMessage deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
+ ObjectMapper mapper = (ObjectMapper) p.getCodec();
+ JsonNode node = mapper.readTree(p);
+
+ String id = getTextOrNull(node, "id");
+ String role = getTextOrNull(node, "role");
+ String toolCallId = getTextOrNull(node, "toolCallId");
+
+ List toolCalls = null;
+ if (node.has("toolCalls") && !node.get("toolCalls").isNull()) {
+ toolCalls =
+ mapper.convertValue(
+ node.get("toolCalls"), new TypeReference>() {});
+ }
+
+ // Handle union type: content can be string or array
+ String content = null;
+ List contentParts = null;
+
+ JsonNode contentNode = node.get("content");
+ if (contentNode != null && !contentNode.isNull()) {
+ if (contentNode.isTextual()) {
+ content = contentNode.asText();
+ } else if (contentNode.isArray()) {
+ contentParts =
+ mapper.convertValue(
+ contentNode, new TypeReference>() {});
+ }
+ }
+
+ return new AguiMessage(id, role, content, contentParts, toolCalls, toolCallId);
+ }
+
+ private String getTextOrNull(JsonNode node, String field) {
+ JsonNode fieldNode = node.get(field);
+ if (fieldNode != null && !fieldNode.isNull()) {
+ return fieldNode.asText();
+ }
+ return null;
+ }
+}
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiTextContent.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiTextContent.java
new file mode 100644
index 000000000..c06650100
--- /dev/null
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiTextContent.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.core.agui.model;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.Objects;
+
+/**
+ * Represents a text content part in the AG-UI protocol's multimodal message format.
+ *
+ *
JSON format: {@code {"type": "text", "text": "..."}}
+ */
+public record AguiTextContent(String text) implements AguiContentPart {
+
+ @JsonCreator
+ public AguiTextContent(@JsonProperty("text") String text) {
+ this.text = Objects.requireNonNull(text, "text cannot be null");
+ }
+}
diff --git a/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiUrlSource.java b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiUrlSource.java
new file mode 100644
index 000000000..bd76ba87b
--- /dev/null
+++ b/agentscope-extensions/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/model/AguiUrlSource.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.core.agui.model;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.Objects;
+
+/**
+ * Represents a URL-based content source in the AG-UI protocol.
+ *
+ *
This configuration scans all Agent beans in the application context and registers them with
- * the registry. It supports both singleton and prototype scoped beans:
+ *
Fixes prototype bean handling: the upstream version uses
+ * {@code beanFactory.getBeansOfType(Agent.class)} which eagerly instantiates
+ * ALL Agent beans including prototypes. This fails when prototype factory methods
+ * depend on request-scoped context (e.g., {@code AguiRequestContext}) that is not
+ * available at application startup.
*
- *
- *
Singleton beans: The bean instance is registered directly
- *
Prototype beans: A factory is registered that creates new instances per request
- * (thread-safe)
- *
- *
- *
Agent ID Resolution:
- *
- *
- *
{@link AguiAgentId} annotation on the bean method or class
- *
This version scans bean definitions without instantiating prototype beans,
+ * registering only a factory supplier for them.
*/
public class AguiAgentAutoRegistration implements BeanFactoryAware, InitializingBean {
@@ -80,11 +48,6 @@ public class AguiAgentAutoRegistration implements BeanFactoryAware, Initializing
private final AguiAgentRegistry registry;
- /**
- * Creates a new AguiAgentAutoRegistration.
- *
- * @param registry The agent registry
- */
public AguiAgentAutoRegistration(AguiAgentRegistry registry) {
this.registry = registry;
}
@@ -99,7 +62,9 @@ public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
/**
* Registers all Agent beans with the AguiAgentRegistry.
*
- *
This method is called after the registry is created and scans for all Agent beans.
+ *
Unlike the upstream version, this does NOT call {@code getBeansOfType(Agent.class)}
+ * which would eagerly instantiate prototype beans. Instead, it iterates over bean
+ * definitions and resolves types without instantiation.
*/
protected void aguiAgentAutoRegistrar() {
if (beanFactory == null) {
@@ -107,11 +72,21 @@ protected void aguiAgentAutoRegistrar() {
return;
}
- Map agentBeans = beanFactory.getBeansOfType(Agent.class);
+ String[] beanNames = beanFactory.getBeanDefinitionNames();
- for (Map.Entry entry : agentBeans.entrySet()) {
- String beanName = entry.getKey();
- Agent agent = entry.getValue();
+ for (String beanName : beanNames) {
+ // Resolve bean type without instantiation
+ Class beanType;
+ try {
+ beanType = beanFactory.getType(beanName);
+ } catch (Exception e) {
+ logger.debug("Could not resolve type for bean '{}': {}", beanName, e.getMessage());
+ continue;
+ }
+
+ if (beanType == null || !Agent.class.isAssignableFrom(beanType)) {
+ continue;
+ }
// Determine agent ID
String agentId = resolveAgentId(beanName);
@@ -122,18 +97,19 @@ protected void aguiAgentAutoRegistrar() {
continue;
}
- // Check bean scope
boolean isPrototype = isPrototypeBean(beanName);
if (isPrototype) {
- // Register factory for prototype beans (thread-safe: new instance per call)
+ // Register factory for prototype beans WITHOUT creating an instance.
+ // The factory will be invoked per-request when the HTTP context is available.
registry.registerFactory(agentId, () -> beanFactory.getBean(beanName, Agent.class));
logger.info(
"Auto-registered prototype agent '{}' (bean: {}) with factory",
agentId,
beanName);
} else {
- // Register singleton directly
+ // Register singleton directly (safe to instantiate at startup)
+ Agent agent = beanFactory.getBean(beanName, Agent.class);
registry.register(agentId, agent);
logger.info("Auto-registered singleton agent '{}' (bean: {})", agentId, beanName);
}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/AguiSessionManager.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/AguiSessionManager.java
new file mode 100644
index 000000000..4fd2a4632
--- /dev/null
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/AguiSessionManager.java
@@ -0,0 +1,108 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.spring.boot.agui.common;
+
+import io.agentscope.core.agent.Agent;
+import java.util.function.Supplier;
+
+/**
+ * Manages agent sessions by threadId for server-side memory management.
+ *
+ *
This interface defines the contract for managing agent instances associated with conversation
+ * threads. Implementations may store sessions in-memory (for single-instance deployments) or
+ * delegate to external stores like Redis (for distributed deployments).
+ *
+ *
Usage:
+ *
+ *
{@code
+ * AguiSessionManager manager = ...;
+ *
+ * // Get or create an agent for a thread
+ * Agent agent = manager.getOrCreateAgent("thread-123", "default", () -> createAgent());
+ *
+ * // Check if agent has memory
+ * boolean hasMemory = manager.hasMemory("thread-123", "default");
+ *
+ * // Clean up expired sessions
+ * manager.cleanupExpiredSessions();
+ * }
+ *
+ * @see InMemoryAguiSessionManager
+ * @see SessionAwareAguiSessionManager
+ */
+public interface AguiSessionManager {
+
+ /**
+ * Get or create an agent for the given threadId.
+ *
+ *
This method should be thread-safe. If an agent already exists for this threadId with the
+ * same agentId, the existing agent should be reused. If the agentId has changed, a new agent
+ * should be created.
+ *
+ * @param threadId The thread identifier
+ * @param agentId The agent type identifier
+ * @param agentFactory Factory to create new agents if needed
+ * @return The agent for this thread
+ */
+ Agent getOrCreateAgent(String threadId, String agentId, Supplier agentFactory);
+
+ /**
+ * Check if a session exists and has memory for the given threadId and agentId.
+ *
+ * @param threadId The thread identifier
+ * @param agentId The agent type identifier
+ * @return true if the session exists and the agent has non-empty memory
+ */
+ boolean hasMemory(String threadId, String agentId);
+
+ /**
+ * Remove a session by threadId and agentId.
+ *
+ * @param threadId The thread identifier
+ * @param agentId The agent type identifier
+ * @return true if a session was removed
+ */
+ boolean removeSession(String threadId, String agentId);
+
+ /** Clean up sessions that have been inactive for longer than the timeout. */
+ void cleanupExpiredSessions();
+
+ /**
+ * Get the current number of active sessions.
+ *
+ * @return Number of sessions
+ */
+ int getSessionCount();
+
+ /**
+ * Save agent state after a request completes.
+ *
+ *
Called when an AG-UI request finishes (the event stream completes).
+ * Implementations that delegate to external session stores should persist
+ * the agent's state here. In-memory implementations can treat this as a no-op
+ * since the agent instance is already cached locally.
+ *
+ * @param threadId The thread identifier
+ * @param agentId The agent type identifier
+ * @param agent The agent whose state should be saved
+ */
+ default void saveAgent(String threadId, String agentId, Agent agent) {
+ // Default no-op: in-memory implementations don't need explicit saves
+ }
+
+ /** Clear all sessions. */
+ void clear();
+}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/DefaultAgentResolver.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/DefaultAgentResolver.java
index 9970f8f9b..666ea8e1d 100644
--- a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/DefaultAgentResolver.java
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/DefaultAgentResolver.java
@@ -19,7 +19,9 @@
import io.agentscope.core.agui.AguiException;
import io.agentscope.core.agui.processor.AgentResolver;
import io.agentscope.core.agui.registry.AguiAgentRegistry;
+import java.util.Map;
import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
/**
* Default implementation of {@link AgentResolver} for Spring Boot integration.
@@ -27,15 +29,18 @@
*
This resolver supports two modes:
*
*
Simple mode: Directly looks up agents from the registry
- *
Session mode: Uses {@link ThreadSessionManager} for server-side memory
+ *
Session mode: Uses {@link AguiSessionManager} for server-side memory
*
*/
public class DefaultAgentResolver implements AgentResolver {
private final AguiAgentRegistry registry;
- private final ThreadSessionManager sessionManager;
+ private final AguiSessionManager sessionManager;
private final boolean serverSideMemory;
+ /** Tracks the agentId used for each threadId during a request lifecycle. */
+ private final Map threadAgentIdMap = new ConcurrentHashMap<>();
+
/**
* Creates a simple resolver without session support.
*
@@ -54,7 +59,7 @@ public DefaultAgentResolver(AguiAgentRegistry registry) {
*/
public DefaultAgentResolver(
AguiAgentRegistry registry,
- ThreadSessionManager sessionManager,
+ AguiSessionManager sessionManager,
boolean serverSideMemory) {
this.registry = Objects.requireNonNull(registry, "registry cannot be null");
this.sessionManager = sessionManager;
@@ -64,6 +69,8 @@ public DefaultAgentResolver(
@Override
public Agent resolveAgent(String agentId, String threadId) {
if (serverSideMemory && sessionManager != null) {
+ // Track the agentId for this threadId so hasMemory/onComplete can use it
+ threadAgentIdMap.put(threadId, agentId);
// Server-side memory mode: use session manager
return sessionManager.getOrCreateAgent(
threadId,
@@ -84,11 +91,23 @@ public Agent resolveAgent(String agentId, String threadId) {
@Override
public boolean hasMemory(String threadId) {
if (serverSideMemory && sessionManager != null) {
- return sessionManager.hasMemory(threadId);
+ String agentId = threadAgentIdMap.get(threadId);
+ return agentId != null && sessionManager.hasMemory(threadId, agentId);
}
return false;
}
+ @Override
+ public void onComplete(String threadId, Agent agent) {
+ if (serverSideMemory && sessionManager != null) {
+ String agentId = threadAgentIdMap.remove(threadId);
+ if (agentId == null) {
+ agentId = agent.getAgentId();
+ }
+ sessionManager.saveAgent(threadId, agentId, agent);
+ }
+ }
+
/**
* Creates a new builder for DefaultAgentResolver.
*
@@ -102,7 +121,7 @@ public static Builder builder() {
public static class Builder {
private AguiAgentRegistry registry;
- private ThreadSessionManager sessionManager;
+ private AguiSessionManager sessionManager;
private boolean serverSideMemory = false;
/**
@@ -122,7 +141,7 @@ public Builder registry(AguiAgentRegistry registry) {
* @param sessionManager The session manager
* @return This builder
*/
- public Builder sessionManager(ThreadSessionManager sessionManager) {
+ public Builder sessionManager(AguiSessionManager sessionManager) {
this.sessionManager = sessionManager;
return this;
}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/ThreadSessionManager.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/InMemoryAguiSessionManager.java
similarity index 65%
rename from agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/ThreadSessionManager.java
rename to agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/InMemoryAguiSessionManager.java
index 3ab8e0350..eccfd7fc0 100644
--- a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/ThreadSessionManager.java
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/InMemoryAguiSessionManager.java
@@ -27,58 +27,53 @@
import org.slf4j.LoggerFactory;
/**
- * Manages agent sessions by threadId for server-side memory management.
+ * In-memory implementation of {@link AguiSessionManager}.
*
- *
This manager maintains a pool of agent instances, each associated with a threadId. When
- * server-side memory is enabled, the same agent instance is reused for requests with the same
- * threadId, preserving conversation history across requests.
+ *
This implementation maintains a pool of agent instances in a {@link ConcurrentHashMap}, each
+ * associated with a threadId. When server-side memory is enabled, the same agent instance is reused
+ * for requests with the same threadId, preserving conversation history across requests.
+ *
+ *
This is suitable for single-instance deployments. For distributed deployments where requests
+ * may be routed to different machines, use {@link SessionAwareAguiSessionManager} instead.
*
*
Usage:
*
*
{@code
- * ThreadSessionManager manager = new ThreadSessionManager(1000, 30);
+ * AguiSessionManager manager = new InMemoryAguiSessionManager(1000, 30);
*
* // Get or create an agent for a thread
* Agent agent = manager.getOrCreateAgent("thread-123", "default", () -> createAgent());
*
* // Check if agent has memory
- * boolean hasMemory = manager.hasMemory("thread-123");
+ * boolean hasMemory = manager.hasMemory("thread-123", "default");
*
* // Clean up expired sessions
* manager.cleanupExpiredSessions();
* }
*/
-public class ThreadSessionManager {
+public class InMemoryAguiSessionManager implements AguiSessionManager {
- private static final Logger logger = LoggerFactory.getLogger(ThreadSessionManager.class);
+ private static final Logger logger = LoggerFactory.getLogger(InMemoryAguiSessionManager.class);
- private final Map sessions = new ConcurrentHashMap<>();
+ private final Map sessions = new ConcurrentHashMap<>();
private final int maxSessions;
private final int sessionTimeoutMinutes;
/**
- * Creates a new ThreadSessionManager.
+ * Creates a new InMemoryAguiSessionManager.
*
* @param maxSessions Maximum number of sessions to maintain
* @param sessionTimeoutMinutes Session timeout in minutes (0 = no timeout)
*/
- public ThreadSessionManager(int maxSessions, int sessionTimeoutMinutes) {
+ public InMemoryAguiSessionManager(int maxSessions, int sessionTimeoutMinutes) {
this.maxSessions = maxSessions;
this.sessionTimeoutMinutes = sessionTimeoutMinutes;
}
- /**
- * Get or create an agent for the given threadId.
- *
- *
This method is thread-safe. It uses atomic operations to ensure that concurrent requests
- * for the same threadId will share the same agent instance.
- *
- * @param threadId The thread identifier
- * @param agentId The agent type identifier
- * @param agentFactory Factory to create new agents if needed
- * @return The agent for this thread
- */
+ @Override
public Agent getOrCreateAgent(String threadId, String agentId, Supplier agentFactory) {
+ String compositeKey = buildCompositeKey(agentId, threadId);
+
// Clean up if we're at capacity
if (sessions.size() >= maxSessions) {
cleanupExpiredSessions();
@@ -89,23 +84,17 @@ public Agent getOrCreateAgent(String threadId, String agentId, Supplier a
}
// Use compute() for atomic check-and-update to avoid race conditions
- ThreadSession session =
+ AguiSession session =
sessions.compute(
- threadId,
+ compositeKey,
(k, existing) -> {
if (existing == null) {
// No existing session, create new one
- logger.debug("Creating new session for threadId: {}", threadId);
- return new ThreadSession(agentId, agentFactory.get());
- }
- if (!existing.getAgentId().equals(agentId)) {
- // Agent type changed, create new session
logger.debug(
- "Agent type changed for threadId {}: {} -> {}",
+ "Creating new session for threadId: {}, agentId: {}",
threadId,
- existing.getAgentId(),
agentId);
- return new ThreadSession(agentId, agentFactory.get());
+ return new AguiSession(agentId, agentFactory.get());
}
// Same agent type, update access time and reuse
existing.updateLastAccess();
@@ -115,14 +104,10 @@ public Agent getOrCreateAgent(String threadId, String agentId, Supplier a
return session.getAgent();
}
- /**
- * Check if a session exists and has memory for the given threadId.
- *
- * @param threadId The thread identifier
- * @return true if the session exists and the agent has non-empty memory
- */
- public boolean hasMemory(String threadId) {
- ThreadSession session = sessions.get(threadId);
+ @Override
+ public boolean hasMemory(String threadId, String agentId) {
+ String compositeKey = buildCompositeKey(agentId, threadId);
+ AguiSession session = sessions.get(compositeKey);
if (session == null) {
return false;
}
@@ -139,26 +124,25 @@ public boolean hasMemory(String threadId) {
}
/**
- * Get the session for a threadId if it exists.
+ * Get the session for a threadId and agentId if it exists.
+ *
+ *
This method is specific to the in-memory implementation and not part of the {@link
+ * AguiSessionManager} interface.
*
* @param threadId The thread identifier
+ * @param agentId The agent type identifier
* @return Optional containing the session, or empty if not found
*/
- public Optional getSession(String threadId) {
- return Optional.ofNullable(sessions.get(threadId));
+ public Optional getSession(String threadId, String agentId) {
+ return Optional.ofNullable(sessions.get(buildCompositeKey(agentId, threadId)));
}
- /**
- * Remove a session by threadId.
- *
- * @param threadId The thread identifier
- * @return true if a session was removed
- */
- public boolean removeSession(String threadId) {
- return sessions.remove(threadId) != null;
+ @Override
+ public boolean removeSession(String threadId, String agentId) {
+ return sessions.remove(buildCompositeKey(agentId, threadId)) != null;
}
- /** Clean up sessions that have been inactive for longer than the timeout. */
+ @Override
public void cleanupExpiredSessions() {
if (sessionTimeoutMinutes <= 0) {
return;
@@ -199,28 +183,24 @@ private void removeOldestSession() {
}
}
- /**
- * Get the current number of active sessions.
- *
- * @return Number of sessions
- */
+ @Override
public int getSessionCount() {
return sessions.size();
}
- /** Clear all sessions. */
+ @Override
public void clear() {
sessions.clear();
}
- /** Represents a thread session with its agent and metadata. */
- public static class ThreadSession {
+ /** Represents a session with its agent and metadata. */
+ public static class AguiSession {
private final String agentId;
private final Agent agent;
private Instant lastAccess;
- ThreadSession(String agentId, Agent agent) {
+ AguiSession(String agentId, Agent agent) {
this.agentId = agentId;
this.agent = agent;
this.lastAccess = Instant.now();
@@ -242,4 +222,8 @@ void updateLastAccess() {
this.lastAccess = Instant.now();
}
}
+
+ private static String buildCompositeKey(String agentId, String threadId) {
+ return agentId + ":" + threadId;
+ }
}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/SessionAwareAguiSessionManager.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/SessionAwareAguiSessionManager.java
new file mode 100644
index 000000000..aa2adb8d7
--- /dev/null
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/common/SessionAwareAguiSessionManager.java
@@ -0,0 +1,155 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.spring.boot.agui.common;
+
+import io.agentscope.core.agent.Agent;
+import io.agentscope.core.session.Session;
+import io.agentscope.core.state.SessionKey;
+import io.agentscope.core.state.SimpleSessionKey;
+import io.agentscope.core.state.StateModule;
+import java.util.Objects;
+import java.util.Set;
+import java.util.function.Supplier;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Distributed implementation of {@link AguiSessionManager} that delegates state management to an
+ * external {@link Session} store (e.g., Redis, MySQL).
+ *
+ *
Unlike {@link InMemoryAguiSessionManager}, this implementation does not maintain a local cache
+ * of agent instances. Each call to {@link #getOrCreateAgent} creates a new agent via the provided
+ * factory and automatically restores its state from the session store via {@code
+ * agent.loadIfExists()}. When a request completes, {@link #saveAgent} persists the agent's state
+ * back via {@code agent.saveTo()}.
+ *
+ *
This is suitable for distributed deployments where requests for the same conversation thread
+ * may be routed to different machines. The shared {@link Session} store ensures conversation state
+ * is accessible from any machine.
+ *
+ *
+ */
+public class SessionAwareAguiSessionManager implements AguiSessionManager {
+
+ private static final Logger logger =
+ LoggerFactory.getLogger(SessionAwareAguiSessionManager.class);
+
+ private final Session session;
+
+ /**
+ * Creates a new SessionAwareAguiSessionManager.
+ *
+ * @param session The shared session store (e.g., RedisSession)
+ */
+ public SessionAwareAguiSessionManager(Session session) {
+ this.session = Objects.requireNonNull(session, "session cannot be null");
+ }
+
+ @Override
+ public Agent getOrCreateAgent(String threadId, String agentId, Supplier agentFactory) {
+ logger.debug("Creating agent for threadId: {}, agentId: {}", threadId, agentId);
+ Agent agent = agentFactory.get();
+
+ // Restore agent state from external session store if the agent supports it
+ if (agent instanceof StateModule stateModule) {
+ SessionKey sessionKey = buildSessionKey(agentId, threadId);
+ boolean loaded = stateModule.loadIfExists(session, sessionKey);
+ logger.debug(
+ "State load for threadId {}, agentId {}: {}",
+ threadId,
+ agentId,
+ loaded ? "restored" : "new session");
+ }
+
+ return agent;
+ }
+
+ @Override
+ public void saveAgent(String threadId, String agentId, Agent agent) {
+ if (agent instanceof StateModule stateModule) {
+ SessionKey sessionKey = buildSessionKey(agentId, threadId);
+ stateModule.saveTo(session, sessionKey);
+ logger.debug("Saved agent state for threadId: {}, agentId: {}", threadId, agentId);
+ }
+ }
+
+ @Override
+ public boolean hasMemory(String threadId, String agentId) {
+ SessionKey sessionKey = buildSessionKey(agentId, threadId);
+ return session.exists(sessionKey);
+ }
+
+ @Override
+ public boolean removeSession(String threadId, String agentId) {
+ SessionKey sessionKey = buildSessionKey(agentId, threadId);
+ if (session.exists(sessionKey)) {
+ session.delete(sessionKey);
+ logger.debug("Removed session for threadId: {}, agentId: {}", threadId, agentId);
+ return true;
+ }
+ return false;
+ }
+
+ @Override
+ public void cleanupExpiredSessions() {
+ // No-op: external session stores (e.g., Redis) manage expiration via TTL
+ }
+
+ @Override
+ public int getSessionCount() {
+ Set keys = session.listSessionKeys();
+ return keys.size();
+ }
+
+ @Override
+ public void clear() {
+ Set keys = session.listSessionKeys();
+ for (SessionKey key : keys) {
+ session.delete(key);
+ }
+ logger.debug("Cleared {} sessions from external store", keys.size());
+ }
+
+ /**
+ * Build a composite session key from agentId and threadId.
+ *
+ *
The resulting key format is {@code agentId:threadId}, ensuring that
+ * different agents maintain separate session state even for the same thread.
+ *
+ * @param agentId the agent type identifier
+ * @param threadId the thread identifier
+ * @return a {@link SessionKey} combining agentId and threadId
+ */
+ private static SessionKey buildSessionKey(String agentId, String threadId) {
+ return SimpleSessionKey.of(agentId + ":" + threadId);
+ }
+}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AgentscopeAguiMvcAutoConfiguration.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AgentscopeAguiMvcAutoConfiguration.java
index e74d5a91b..9e9177b54 100644
--- a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AgentscopeAguiMvcAutoConfiguration.java
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AgentscopeAguiMvcAutoConfiguration.java
@@ -19,7 +19,8 @@
import io.agentscope.core.agui.adapter.AguiAdapterConfig;
import io.agentscope.core.agui.registry.AguiAgentRegistry;
import io.agentscope.spring.boot.agui.common.AguiProperties;
-import io.agentscope.spring.boot.agui.common.ThreadSessionManager;
+import io.agentscope.spring.boot.agui.common.AguiSessionManager;
+import io.agentscope.spring.boot.agui.common.InMemoryAguiSessionManager;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -50,15 +51,15 @@
public class AgentscopeAguiMvcAutoConfiguration {
/**
- * Creates the thread session manager bean.
+ * Creates the session manager bean.
*
* @param props The configuration properties
- * @return A new ThreadSessionManager
+ * @return A new InMemoryAguiSessionManager
*/
@Bean
@ConditionalOnMissingBean
- public ThreadSessionManager threadSessionManager(AguiProperties props) {
- return new ThreadSessionManager(
+ public AguiSessionManager aguiSessionManager(AguiProperties props) {
+ return new InMemoryAguiSessionManager(
props.getMaxThreadSessions(), props.getSessionTimeoutMinutes());
}
@@ -66,14 +67,14 @@ public ThreadSessionManager threadSessionManager(AguiProperties props) {
* Creates the AG-UI MVC controller bean.
*
* @param registry The agent registry
- * @param sessionManager The thread session manager
+ * @param sessionManager The session manager
* @param props The configuration properties
* @return A new AguiMvcController
*/
@Bean
@ConditionalOnMissingBean
public AguiMvcController aguiMvcController(
- AguiAgentRegistry registry, ThreadSessionManager sessionManager, AguiProperties props) {
+ AguiAgentRegistry registry, AguiSessionManager sessionManager, AguiProperties props) {
AguiAdapterConfig config =
AguiAdapterConfig.builder()
.toolMergeMode(props.getDefaultToolMergeMode())
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AguiMvcController.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AguiMvcController.java
index 52f2ca595..5f1a523eb 100644
--- a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AguiMvcController.java
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AguiMvcController.java
@@ -16,15 +16,18 @@
package io.agentscope.spring.boot.agui.mvc;
import io.agentscope.core.agui.AguiException;
+import io.agentscope.core.agui.AguiRequestContext;
import io.agentscope.core.agui.adapter.AguiAdapterConfig;
import io.agentscope.core.agui.encoder.AguiEventEncoder;
import io.agentscope.core.agui.event.AguiEvent;
import io.agentscope.core.agui.model.RunAgentInput;
import io.agentscope.core.agui.processor.AguiRequestProcessor;
import io.agentscope.core.agui.registry.AguiAgentRegistry;
+import io.agentscope.spring.boot.agui.common.AguiSessionManager;
import io.agentscope.spring.boot.agui.common.DefaultAgentResolver;
-import io.agentscope.spring.boot.agui.common.ThreadSessionManager;
import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -99,7 +102,25 @@ private AguiMvcController(Builder builder) {
* @return An SseEmitter for streaming AG-UI events
*/
public SseEmitter handle(RunAgentInput input, String headerAgentId) {
- return handleInternal(input, headerAgentId, null);
+ return handleInternal(
+ input, headerAgentId, null, Collections.emptyMap(), Collections.emptyMap());
+ }
+
+ /**
+ * Handle an AG-UI run request with HTTP request headers and parameters.
+ *
+ * @param input The run agent input
+ * @param headerAgentId The agent ID from HTTP header (may be null)
+ * @param requestHeaders All HTTP request headers
+ * @param requestParams All HTTP query parameters
+ * @return An SseEmitter for streaming AG-UI events
+ */
+ public SseEmitter handle(
+ RunAgentInput input,
+ String headerAgentId,
+ Map> requestHeaders,
+ Map> requestParams) {
+ return handleInternal(input, headerAgentId, null, requestHeaders, requestParams);
}
/**
@@ -112,11 +133,36 @@ public SseEmitter handle(RunAgentInput input, String headerAgentId) {
*/
public SseEmitter handleWithAgentId(
RunAgentInput input, String headerAgentId, String pathAgentId) {
- return handleInternal(input, headerAgentId, pathAgentId);
+ return handleInternal(
+ input, headerAgentId, pathAgentId, Collections.emptyMap(), Collections.emptyMap());
+ }
+
+ /**
+ * Handle an AG-UI run request with agent ID in the URL path,
+ * plus HTTP request headers and parameters.
+ *
+ * @param input The run agent input
+ * @param headerAgentId The agent ID from HTTP header (may be null)
+ * @param pathAgentId The agent ID from URL path variable
+ * @param requestHeaders All HTTP request headers
+ * @param requestParams All HTTP query parameters
+ * @return An SseEmitter for streaming AG-UI events
+ */
+ public SseEmitter handleWithAgentId(
+ RunAgentInput input,
+ String headerAgentId,
+ String pathAgentId,
+ Map> requestHeaders,
+ Map> requestParams) {
+ return handleInternal(input, headerAgentId, pathAgentId, requestHeaders, requestParams);
}
private SseEmitter handleInternal(
- RunAgentInput input, String headerAgentId, String pathAgentId) {
+ RunAgentInput input,
+ String headerAgentId,
+ String pathAgentId,
+ Map> requestHeaders,
+ Map> requestParams) {
SseEmitter emitter = new SseEmitter(sseTimeout);
String threadId = input.getThreadId();
String runId = input.getRunId();
@@ -125,6 +171,7 @@ private SseEmitter handleInternal(
() -> {
Disposable subscription = null;
try {
+ AguiRequestContext.init(requestHeaders, requestParams);
// Process request - returns both agent and event stream
AguiRequestProcessor.ProcessResult result =
processor.process(input, headerAgentId, pathAgentId);
@@ -182,6 +229,8 @@ private SseEmitter handleInternal(
} catch (Exception e) {
logger.error("Error processing AG-UI request: {}", e.getMessage());
sendErrorAndComplete(emitter, threadId, runId, e.getMessage());
+ } finally {
+ AguiRequestContext.clear();
}
});
@@ -239,7 +288,7 @@ public static Builder builder() {
public static class Builder {
private AguiAgentRegistry registry;
- private ThreadSessionManager sessionManager;
+ private AguiSessionManager sessionManager;
private AguiAdapterConfig config;
private boolean serverSideMemory = false;
private String agentIdHeader;
@@ -257,12 +306,12 @@ public Builder agentRegistry(AguiAgentRegistry registry) {
}
/**
- * Set the thread session manager for server-side memory support.
+ * Set the session manager for server-side memory support.
*
* @param sessionManager The session manager
* @return This builder
*/
- public Builder sessionManager(ThreadSessionManager sessionManager) {
+ public Builder sessionManager(AguiSessionManager sessionManager) {
this.sessionManager = sessionManager;
return this;
}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AguiRestController.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AguiRestController.java
index dfaa0902a..979f0785d 100644
--- a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AguiRestController.java
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/mvc/AguiRestController.java
@@ -16,6 +16,13 @@
package io.agentscope.spring.boot.agui.mvc;
import io.agentscope.core.agui.model.RunAgentInput;
+import jakarta.servlet.http.HttpServletRequest;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
@@ -29,6 +36,10 @@
*
*
This controller exposes the AG-UI run endpoints for Spring MVC applications.
* It delegates the actual processing to {@link AguiMvcController}.
+ *
+ *
This version extracts all HTTP request headers and query parameters from the
+ * {@link HttpServletRequest} and passes them to the controller for propagation
+ * via {@link io.agentscope.core.agui.AguiRequestContext}.
*/
@RestController
public class AguiRestController {
@@ -62,6 +73,7 @@ public AguiRestController(
*
"default"
*
*
+ * @param request The HTTP servlet request
* @param input The run agent input
* @param agentIdHeader The agent ID from HTTP header (optional)
* @return An SseEmitter for streaming AG-UI events
@@ -71,12 +83,15 @@ public AguiRestController(
consumes = MediaType.APPLICATION_JSON_VALUE,
produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public SseEmitter run(
+ HttpServletRequest request,
@RequestBody RunAgentInput input,
@RequestHeader(
value = "${agentscope.agui.agent-id-header:X-Agent-Id}",
required = false)
String agentIdHeader) {
- return aguiMvcController.handle(input, agentIdHeader);
+ Map> headers = extractHeaders(request);
+ Map> params = extractParameters(request);
+ return aguiMvcController.handle(input, agentIdHeader, headers, params);
}
/**
@@ -85,6 +100,7 @@ public SseEmitter run(
*
The path variable takes highest priority for agent resolution.
*
* @param agentId The agent ID from path variable
+ * @param request The HTTP servlet request
* @param input The run agent input
* @param agentIdHeader The agent ID from HTTP header (optional)
* @return An SseEmitter for streaming AG-UI events
@@ -94,12 +110,33 @@ public SseEmitter run(
consumes = MediaType.APPLICATION_JSON_VALUE,
produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public SseEmitter runWithAgentId(
- @PathVariable String agentId,
+ @PathVariable("agentId") String agentId,
+ HttpServletRequest request,
@RequestBody RunAgentInput input,
@RequestHeader(
value = "${agentscope.agui.agent-id-header:X-Agent-Id}",
required = false)
String agentIdHeader) {
- return aguiMvcController.handleWithAgentId(input, agentIdHeader, agentId);
+ Map> headers = extractHeaders(request);
+ Map> params = extractParameters(request);
+ return aguiMvcController.handleWithAgentId(input, agentIdHeader, agentId, headers, params);
+ }
+
+ private static Map> extractHeaders(HttpServletRequest request) {
+ Map> headers = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
+ Enumeration headerNames = request.getHeaderNames();
+ if (headerNames != null) {
+ while (headerNames.hasMoreElements()) {
+ String name = headerNames.nextElement();
+ headers.put(name, Collections.list(request.getHeaders(name)));
+ }
+ }
+ return headers;
+ }
+
+ private static Map> extractParameters(HttpServletRequest request) {
+ Map> params = new LinkedHashMap<>();
+ request.getParameterMap().forEach((name, values) -> params.put(name, List.of(values)));
+ return params;
}
}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/webflux/AgentscopeAguiWebFluxAutoConfiguration.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/webflux/AgentscopeAguiWebFluxAutoConfiguration.java
index 23513f2b4..ce827f3a9 100644
--- a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/webflux/AgentscopeAguiWebFluxAutoConfiguration.java
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/webflux/AgentscopeAguiWebFluxAutoConfiguration.java
@@ -19,7 +19,8 @@
import io.agentscope.core.agui.adapter.AguiAdapterConfig;
import io.agentscope.core.agui.registry.AguiAgentRegistry;
import io.agentscope.spring.boot.agui.common.AguiProperties;
-import io.agentscope.spring.boot.agui.common.ThreadSessionManager;
+import io.agentscope.spring.boot.agui.common.AguiSessionManager;
+import io.agentscope.spring.boot.agui.common.InMemoryAguiSessionManager;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -53,15 +54,15 @@
public class AgentscopeAguiWebFluxAutoConfiguration {
/**
- * Creates the thread session manager bean.
+ * Creates the session manager bean.
*
* @param props The configuration properties
- * @return A new ThreadSessionManager
+ * @return A new InMemoryAguiSessionManager
*/
@Bean
@ConditionalOnMissingBean
- public ThreadSessionManager threadSessionManager(AguiProperties props) {
- return new ThreadSessionManager(
+ public AguiSessionManager aguiSessionManager(AguiProperties props) {
+ return new InMemoryAguiSessionManager(
props.getMaxThreadSessions(), props.getSessionTimeoutMinutes());
}
@@ -69,14 +70,14 @@ public ThreadSessionManager threadSessionManager(AguiProperties props) {
* Creates the AG-UI WebFlux handler bean.
*
* @param registry The agent registry
- * @param sessionManager The thread session manager
+ * @param sessionManager The session manager
* @param props The configuration properties
* @return A new AguiWebFluxHandler
*/
@Bean
@ConditionalOnMissingBean
public AguiWebFluxHandler aguiWebFluxHandler(
- AguiAgentRegistry registry, ThreadSessionManager sessionManager, AguiProperties props) {
+ AguiAgentRegistry registry, AguiSessionManager sessionManager, AguiProperties props) {
AguiAdapterConfig config =
AguiAdapterConfig.builder()
.toolMergeMode(props.getDefaultToolMergeMode())
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/webflux/AguiWebFluxHandler.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/webflux/AguiWebFluxHandler.java
index 50f07f647..e46c36227 100644
--- a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/webflux/AguiWebFluxHandler.java
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/java/io/agentscope/spring/boot/agui/webflux/AguiWebFluxHandler.java
@@ -16,15 +16,19 @@
package io.agentscope.spring.boot.agui.webflux;
import io.agentscope.core.agui.AguiException;
+import io.agentscope.core.agui.AguiRequestContext;
import io.agentscope.core.agui.adapter.AguiAdapterConfig;
import io.agentscope.core.agui.encoder.AguiEventEncoder;
import io.agentscope.core.agui.event.AguiEvent;
import io.agentscope.core.agui.model.RunAgentInput;
import io.agentscope.core.agui.processor.AguiRequestProcessor;
import io.agentscope.core.agui.registry.AguiAgentRegistry;
+import io.agentscope.spring.boot.agui.common.AguiSessionManager;
import io.agentscope.spring.boot.agui.common.DefaultAgentResolver;
-import io.agentscope.spring.boot.agui.common.ThreadSessionManager;
+import java.util.LinkedHashMap;
+import java.util.List;
import java.util.Map;
+import java.util.TreeMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
@@ -129,7 +133,14 @@ private Mono processInput(
String threadId = input.getThreadId();
String runId = input.getRunId();
+ // Extract headers and query parameters for AguiRequestContext
+ Map> headers = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
+ request.headers().asHttpHeaders().forEach(headers::put);
+ Map> params = new LinkedHashMap<>(request.queryParams());
+
try {
+ AguiRequestContext.init(headers, params);
+
// Get header agent ID
String headerAgentId = request.headers().firstHeader(agentIdHeader);
@@ -165,6 +176,8 @@ private Mono processInput(
} catch (Exception e) {
logger.error("Error processing AG-UI request: {}", e.getMessage());
return createErrorResponse(threadId, runId, e.getMessage());
+ } finally {
+ AguiRequestContext.clear();
}
}
@@ -221,7 +234,7 @@ public static Builder builder() {
public static class Builder {
private AguiAgentRegistry registry;
- private ThreadSessionManager sessionManager;
+ private AguiSessionManager sessionManager;
private AguiAdapterConfig config;
private boolean serverSideMemory = false;
private String agentIdHeader;
@@ -238,12 +251,12 @@ public Builder agentRegistry(AguiAgentRegistry registry) {
}
/**
- * Set the thread session manager for server-side memory support.
+ * Set the session manager for server-side memory support.
*
* @param sessionManager The session manager
* @return This builder
*/
- public Builder sessionManager(ThreadSessionManager sessionManager) {
+ public Builder sessionManager(AguiSessionManager sessionManager) {
this.sessionManager = sessionManager;
return this;
}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/resources/META-INF/spring.factories b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/resources/META-INF/spring.factories
new file mode 100644
index 000000000..59cd6f9f7
--- /dev/null
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/main/resources/META-INF/spring.factories
@@ -0,0 +1,4 @@
+org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
+io.agentscope.spring.boot.agui.common.AguiAgentRegistryAutoConfiguration,\
+io.agentscope.spring.boot.agui.mvc.AgentscopeAguiMvcAutoConfiguration,\
+io.agentscope.spring.boot.agui.webflux.AgentscopeAguiWebFluxAutoConfiguration
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/common/InMemoryAguiSessionManagerTest.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/common/InMemoryAguiSessionManagerTest.java
new file mode 100644
index 000000000..ddb3d1da0
--- /dev/null
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/common/InMemoryAguiSessionManagerTest.java
@@ -0,0 +1,307 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.spring.boot.agui.common;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+
+import io.agentscope.core.agent.Agent;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Unit tests for {@link InMemoryAguiSessionManager}.
+ *
+ *
Covers the new composite-key (agentId:threadId) isolation behavior and
+ * the updated API signatures: {@code hasMemory(threadId, agentId)},
+ * {@code removeSession(threadId, agentId)}, and {@code getSession(threadId, agentId)}.
+ */
+@Tag("unit")
+@DisplayName("InMemoryAguiSessionManager Unit Tests")
+class InMemoryAguiSessionManagerTest {
+
+ private InMemoryAguiSessionManager manager;
+
+ @BeforeEach
+ void setUp() {
+ manager = new InMemoryAguiSessionManager(10, 30);
+ }
+
+ // ---- getOrCreateAgent ----
+
+ @Test
+ @DisplayName("Should create new agent for new threadId")
+ void testGetOrCreateAgent_newThread() {
+ Agent mockAgent = mock(Agent.class);
+
+ Agent result = manager.getOrCreateAgent("thread-1", "agent-1", () -> mockAgent);
+
+ assertSame(mockAgent, result);
+ assertEquals(1, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName("Should reuse existing agent for same threadId and agentId")
+ void testGetOrCreateAgent_sameThreadSameAgent() {
+ Agent mockAgent = mock(Agent.class);
+
+ Agent first = manager.getOrCreateAgent("thread-1", "agent-1", () -> mockAgent);
+ Agent second = manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ assertSame(first, second);
+ assertEquals(1, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName(
+ "Should create SEPARATE sessions for same threadId with different agentId (composite"
+ + " key isolation)")
+ void testGetOrCreateAgent_sameThreadDifferentAgent_createsIsolatedSessions() {
+ Agent agent1 = mock(Agent.class);
+ Agent agent2 = mock(Agent.class);
+
+ Agent first = manager.getOrCreateAgent("thread-1", "agent-1", () -> agent1);
+ Agent second = manager.getOrCreateAgent("thread-1", "agent-2", () -> agent2);
+
+ // Different agentIds → different composite keys → both sessions exist independently
+ assertSame(agent1, first);
+ assertSame(agent2, second);
+ assertNotSame(first, second);
+ // Both sessions coexist (agent-1:thread-1 and agent-2:thread-1)
+ assertEquals(2, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName("Should manage multiple threads independently")
+ void testGetOrCreateAgent_multipleThreads() {
+ Agent agent1 = mock(Agent.class);
+ Agent agent2 = mock(Agent.class);
+
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> agent1);
+ manager.getOrCreateAgent("thread-2", "agent-1", () -> agent2);
+
+ assertEquals(2, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName("Same agentId with different threadIds are independent sessions")
+ void testGetOrCreateAgent_differentThreadsSameAgent() {
+ Agent agent1 = mock(Agent.class);
+ Agent agent2 = mock(Agent.class);
+
+ Agent result1 = manager.getOrCreateAgent("thread-A", "chatAgent", () -> agent1);
+ Agent result2 = manager.getOrCreateAgent("thread-B", "chatAgent", () -> agent2);
+
+ assertSame(agent1, result1);
+ assertSame(agent2, result2);
+ assertNotSame(result1, result2);
+ assertEquals(2, manager.getSessionCount());
+ }
+
+ // ---- removeSession(threadId, agentId) — new API ----
+
+ @Test
+ @DisplayName("Should remove session by threadId and agentId")
+ void testRemoveSession_withAgentId() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ assertTrue(manager.removeSession("thread-1", "agent-1"));
+ assertEquals(0, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName("Should return false when removing non-existent session (wrong agentId)")
+ void testRemoveSession_wrongAgentId() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ // Correct threadId but wrong agentId → should NOT remove the existing session
+ assertFalse(manager.removeSession("thread-1", "agent-wrong"));
+ assertEquals(1, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName("Should return false when removing non-existent session (wrong threadId)")
+ void testRemoveSession_wrongThreadId() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ assertFalse(manager.removeSession("thread-wrong", "agent-1"));
+ assertEquals(1, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName("Should return false when removing non-existent session")
+ void testRemoveSession_nonExistent() {
+ assertFalse(manager.removeSession("non-existent", "agent-1"));
+ }
+
+ @Test
+ @DisplayName("Removing one composite key should not affect other sessions for same thread")
+ void testRemoveSession_onlyRemovesMatchingCompositeKey() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+ manager.getOrCreateAgent("thread-1", "agent-2", () -> mock(Agent.class));
+
+ // Only remove agent-1:thread-1, agent-2:thread-1 should remain
+ assertTrue(manager.removeSession("thread-1", "agent-1"));
+ assertEquals(1, manager.getSessionCount());
+ assertTrue(manager.getSession("thread-1", "agent-2").isPresent());
+ }
+
+ // ---- hasMemory(threadId, agentId) — new API ----
+
+ @Test
+ @DisplayName("Should return false for hasMemory on non-existent thread+agentId")
+ void testHasMemory_nonExistentSession() {
+ assertFalse(manager.hasMemory("non-existent", "agent-1"));
+ }
+
+ @Test
+ @DisplayName("Should return false for hasMemory with wrong agentId")
+ void testHasMemory_wrongAgentId() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ assertFalse(manager.hasMemory("thread-1", "agent-wrong"));
+ }
+
+ @Test
+ @DisplayName("Should return false for hasMemory on non-ReActAgent")
+ void testHasMemory_nonReActAgent() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ // Non-ReActAgent → memory check returns false
+ assertFalse(manager.hasMemory("thread-1", "agent-1"));
+ }
+
+ // ---- getSession(threadId, agentId) — new API ----
+
+ @Test
+ @DisplayName("getSession should return session for existing threadId and agentId")
+ void testGetSession_existing() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ assertTrue(manager.getSession("thread-1", "agent-1").isPresent());
+ }
+
+ @Test
+ @DisplayName("getSession should return empty for wrong agentId")
+ void testGetSession_wrongAgentId() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ assertFalse(manager.getSession("thread-1", "agent-wrong").isPresent());
+ }
+
+ @Test
+ @DisplayName("getSession should return empty for non-existent threadId")
+ void testGetSession_nonExistent() {
+ assertFalse(manager.getSession("non-existent", "agent-1").isPresent());
+ }
+
+ @Test
+ @DisplayName(
+ "getSession should independently access different agentId sessions for same thread")
+ void testGetSession_multipleAgentsPerThread() {
+ Agent agentA = mock(Agent.class);
+ Agent agentB = mock(Agent.class);
+ manager.getOrCreateAgent("thread-1", "agent-a", () -> agentA);
+ manager.getOrCreateAgent("thread-1", "agent-b", () -> agentB);
+
+ assertTrue(manager.getSession("thread-1", "agent-a").isPresent());
+ assertTrue(manager.getSession("thread-1", "agent-b").isPresent());
+ assertSame(agentA, manager.getSession("thread-1", "agent-a").get().getAgent());
+ assertSame(agentB, manager.getSession("thread-1", "agent-b").get().getAgent());
+ }
+
+ // ---- clear ----
+
+ @Test
+ @DisplayName("Should clear all sessions")
+ void testClear() {
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+ manager.getOrCreateAgent("thread-2", "agent-1", () -> mock(Agent.class));
+ manager.getOrCreateAgent("thread-1", "agent-2", () -> mock(Agent.class));
+
+ manager.clear();
+
+ assertEquals(0, manager.getSessionCount());
+ }
+
+ // ---- Capacity eviction ----
+
+ @Test
+ @DisplayName("Should evict oldest session when at max capacity")
+ void testMaxSessionsEviction() {
+ InMemoryAguiSessionManager smallManager = new InMemoryAguiSessionManager(2, 30);
+
+ smallManager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+ smallManager.getOrCreateAgent("thread-2", "agent-1", () -> mock(Agent.class));
+ // This should trigger eviction
+ smallManager.getOrCreateAgent("thread-3", "agent-1", () -> mock(Agent.class));
+
+ assertEquals(2, smallManager.getSessionCount());
+ }
+
+ // ---- cleanupExpiredSessions ----
+
+ @Test
+ @DisplayName("cleanupExpiredSessions should be a no-op when timeout is 0")
+ void testCleanupExpiredSessions_noTimeout() {
+ InMemoryAguiSessionManager noTimeoutManager = new InMemoryAguiSessionManager(10, 0);
+ noTimeoutManager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ noTimeoutManager.cleanupExpiredSessions();
+
+ assertEquals(1, noTimeoutManager.getSessionCount());
+ }
+
+ // ---- AguiSession inner class ----
+
+ @Test
+ @DisplayName("AguiSession should expose agentId and agent")
+ void testAguiSession_fields() {
+ Agent mockAgent = mock(Agent.class);
+ manager.getOrCreateAgent("thread-1", "my-agent", () -> mockAgent);
+
+ InMemoryAguiSessionManager.AguiSession session =
+ manager.getSession("thread-1", "my-agent").orElseThrow();
+
+ assertEquals("my-agent", session.getAgentId());
+ assertSame(mockAgent, session.getAgent());
+ }
+
+ @Test
+ @DisplayName("AguiSession lastAccess should update on re-access")
+ void testAguiSession_lastAccessUpdated() throws InterruptedException {
+ Agent mockAgent = mock(Agent.class);
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mockAgent);
+
+ InMemoryAguiSessionManager.AguiSession session =
+ manager.getSession("thread-1", "agent-1").orElseThrow();
+ java.time.Instant firstAccess = session.getLastAccess();
+
+ // Wait a bit then re-access
+ Thread.sleep(10);
+ manager.getOrCreateAgent("thread-1", "agent-1", () -> mock(Agent.class));
+
+ java.time.Instant secondAccess = session.getLastAccess();
+ // Last access should have been updated
+ assertTrue(secondAccess.isAfter(firstAccess) || secondAccess.equals(firstAccess));
+ }
+}
diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/common/SessionAwareAguiSessionManagerTest.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/common/SessionAwareAguiSessionManagerTest.java
new file mode 100644
index 000000000..7d69b34e1
--- /dev/null
+++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/common/SessionAwareAguiSessionManagerTest.java
@@ -0,0 +1,219 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.agentscope.spring.boot.agui.common;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import io.agentscope.core.agent.Agent;
+import io.agentscope.core.session.Session;
+import io.agentscope.core.state.SessionKey;
+import io.agentscope.core.state.SimpleSessionKey;
+import io.agentscope.core.state.StateModule;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Unit tests for {@link SessionAwareAguiSessionManager}.
+ */
+@Tag("unit")
+@DisplayName("SessionAwareAguiSessionManager Unit Tests")
+class SessionAwareAguiSessionManagerTest {
+
+ private Session mockSession;
+ private SessionAwareAguiSessionManager manager;
+
+ @BeforeEach
+ void setUp() {
+ mockSession = mock(Session.class);
+ manager = new SessionAwareAguiSessionManager(mockSession);
+ }
+
+ @Test
+ @DisplayName("Should throw NullPointerException when session is null")
+ void testConstructor_nullSession() {
+ assertThrows(NullPointerException.class, () -> new SessionAwareAguiSessionManager(null));
+ }
+
+ @Test
+ @DisplayName("Should always create new agent via factory")
+ void testGetOrCreateAgent_alwaysCreatesNew() {
+ Agent agent1 = mock(Agent.class);
+ Agent agent2 = mock(Agent.class);
+
+ Agent first = manager.getOrCreateAgent("thread-1", "agent-1", () -> agent1);
+ Agent second = manager.getOrCreateAgent("thread-1", "agent-1", () -> agent2);
+
+ assertSame(agent1, first);
+ assertSame(agent2, second);
+ assertNotSame(first, second);
+ }
+
+ @Test
+ @DisplayName("Should call loadIfExists on StateModule agents during getOrCreateAgent")
+ void testGetOrCreateAgent_loadsStateForStatefulAgent() {
+ StatefulAgent statefulAgent = mock(StatefulAgent.class);
+ when(statefulAgent.loadIfExists(any(Session.class), any(SessionKey.class)))
+ .thenReturn(true);
+
+ Agent result = manager.getOrCreateAgent("thread-1", "agent-1", () -> statefulAgent);
+
+ assertSame(statefulAgent, result);
+ // SessionAwareAguiSessionManager uses composite key "agentId:threadId"
+ verify(statefulAgent).loadIfExists(mockSession, SimpleSessionKey.of("agent-1:thread-1"));
+ }
+
+ @Test
+ @DisplayName("Should not call loadIfExists on plain Agent (non-StateModule)")
+ void testGetOrCreateAgent_skipsLoadForPlainAgent() {
+ Agent plainAgent = mock(Agent.class);
+
+ Agent result = manager.getOrCreateAgent("thread-1", "agent-1", () -> plainAgent);
+
+ assertSame(plainAgent, result);
+ // No loadIfExists call since Agent doesn't implement StateModule
+ }
+
+ @Test
+ @DisplayName("Should call saveTo on StateModule agents via saveAgent")
+ void testSaveAgent_savesStateForStatefulAgent() {
+ StatefulAgent statefulAgent = mock(StatefulAgent.class);
+
+ manager.saveAgent("thread-1", "agent-1", statefulAgent);
+
+ // SessionAwareAguiSessionManager uses composite key "agentId:threadId"
+ verify(statefulAgent).saveTo(mockSession, SimpleSessionKey.of("agent-1:thread-1"));
+ }
+
+ @Test
+ @DisplayName("Should be a no-op for plain Agent (non-StateModule) via saveAgent")
+ void testSaveAgent_skipsForPlainAgent() {
+ Agent plainAgent = mock(Agent.class);
+
+ manager.saveAgent("thread-1", "agent-1", plainAgent);
+
+ // No saveTo call since Agent doesn't implement StateModule
+ }
+
+ @Test
+ @DisplayName("Should delegate hasMemory to session.exists using composite key")
+ void testHasMemory_exists() {
+ // SessionAwareAguiSessionManager uses composite key "agentId:threadId"
+ when(mockSession.exists(SimpleSessionKey.of("agent-1:thread-1"))).thenReturn(true);
+
+ assertTrue(manager.hasMemory("thread-1", "agent-1"));
+ verify(mockSession).exists(SimpleSessionKey.of("agent-1:thread-1"));
+ }
+
+ @Test
+ @DisplayName("Should return false for hasMemory when session does not exist")
+ void testHasMemory_notExists() {
+ when(mockSession.exists(SimpleSessionKey.of("agent-1:thread-1"))).thenReturn(false);
+
+ assertFalse(manager.hasMemory("thread-1", "agent-1"));
+ }
+
+ @Test
+ @DisplayName("Should delete session when it exists using composite key")
+ void testRemoveSession_exists() {
+ // SessionAwareAguiSessionManager uses composite key "agentId:threadId"
+ when(mockSession.exists(SimpleSessionKey.of("agent-1:thread-1"))).thenReturn(true);
+
+ assertTrue(manager.removeSession("thread-1", "agent-1"));
+ verify(mockSession).delete(SimpleSessionKey.of("agent-1:thread-1"));
+ }
+
+ @Test
+ @DisplayName("Should return false when removing non-existent session")
+ void testRemoveSession_notExists() {
+ when(mockSession.exists(SimpleSessionKey.of("agent-1:thread-1"))).thenReturn(false);
+
+ assertFalse(manager.removeSession("thread-1", "agent-1"));
+ verify(mockSession, never()).delete(any());
+ }
+
+ @Test
+ @DisplayName("cleanupExpiredSessions should be a no-op")
+ void testCleanupExpiredSessions() {
+ manager.cleanupExpiredSessions();
+
+ // Should not interact with session at all
+ verify(mockSession, never()).delete(any());
+ verify(mockSession, never()).exists(any());
+ }
+
+ @Test
+ @DisplayName("Should return count from session.listSessionKeys")
+ void testGetSessionCount() {
+ Set keys = new HashSet<>();
+ keys.add(SimpleSessionKey.of("thread-1"));
+ keys.add(SimpleSessionKey.of("thread-2"));
+ when(mockSession.listSessionKeys()).thenReturn(keys);
+
+ assertEquals(2, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName("Should return 0 when no sessions exist")
+ void testGetSessionCount_empty() {
+ when(mockSession.listSessionKeys()).thenReturn(Collections.emptySet());
+
+ assertEquals(0, manager.getSessionCount());
+ }
+
+ @Test
+ @DisplayName("Should delete all session keys on clear")
+ void testClear() {
+ Set keys = new HashSet<>();
+ keys.add(SimpleSessionKey.of("thread-1"));
+ keys.add(SimpleSessionKey.of("thread-2"));
+ when(mockSession.listSessionKeys()).thenReturn(keys);
+
+ manager.clear();
+
+ verify(mockSession).delete(SimpleSessionKey.of("thread-1"));
+ verify(mockSession).delete(SimpleSessionKey.of("thread-2"));
+ verify(mockSession, times(2)).delete(any());
+ }
+
+ @Test
+ @DisplayName("Should handle empty keys on clear")
+ void testClear_empty() {
+ when(mockSession.listSessionKeys()).thenReturn(Collections.emptySet());
+
+ manager.clear();
+
+ verify(mockSession, never()).delete(any());
+ }
+
+ /** Helper interface for mocking an Agent that also implements StateModule. */
+ interface StatefulAgent extends Agent, StateModule {}
+}