Skip to content

Commit b51693f

Browse files
jujnjujn
committed
refactor
1 parent b566723 commit b51693f

6 files changed

Lines changed: 319 additions & 137 deletions

File tree

‎agentscope-core/src/main/java/io/agentscope/core/tool/subagent/SubAgentConfig.java‎

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,10 @@
2222
import java.util.ArrayList;
2323
import java.util.Collections;
2424
import java.util.HashMap;
25+
import java.util.LinkedHashSet;
2526
import java.util.List;
2627
import java.util.Map;
28+
import java.util.Set;
2729

2830
/**
2931
* Configuration for sub-agent registration.
@@ -40,8 +42,16 @@
4042
* an existing one.
4143
* </ul>
4244
*
43-
* <p>Users can also define additional custom parameters to be passed to the sub-agent via
44-
* the {@link Builder#addParameter} methods.
45+
* <p>Users can also define additional parameters to be passed to the sub-agent. The framework
46+
* strictly separates these into two categories for security:
47+
*
48+
* <ul>
49+
* <li><b>Custom Parameters:</b> Added via {@link Builder#addParameter}. These are exposed to the
50+
* LLM in the tool schema, allowing the LLM to infer and generate values based on the conversation.
51+
* <li><b>System Parameters:</b> Added via {@link Builder#addSystemParameter}. These are strictly
52+
* injected by the system via {@code ToolExecutionContext}. They are completely transparent
53+
* (invisible) to the LLM, preventing prompt injection.
54+
* </ul>
4555
*
4656
* <p><b>Default Behavior:</b>
4757
*
@@ -76,6 +86,7 @@ public class SubAgentConfig {
7686
private final Session session;
7787
private final Map<String, Map<String, Object>> customParameters;
7888
private final List<String> requiredCustomParameters;
89+
private final Set<String> systemParameters;
7990

8091
private SubAgentConfig(Builder builder) {
8192
this.toolName = builder.toolName;
@@ -89,6 +100,8 @@ private SubAgentConfig(Builder builder) {
89100
builder.requiredCustomParameters != null
90101
? builder.requiredCustomParameters
91102
: new ArrayList<>();
103+
this.systemParameters =
104+
builder.systemParameters != null ? builder.systemParameters : new LinkedHashSet<>();
92105
}
93106

94107
/**
@@ -187,6 +200,20 @@ public List<String> getRequiredCustomParameters() {
187200
: Collections.unmodifiableList(requiredCustomParameters);
188201
}
189202

203+
/**
204+
* Gets the system parameters defined for the sub-agent tool.
205+
*
206+
* <p>System parameters are injected transparently via ToolExecutionContext
207+
* and are NOT exposed to the LLM in the JSON schema.
208+
*
209+
* @return A set containing the names of system parameters
210+
*/
211+
public Set<String> getSystemParameters() {
212+
return systemParameters == null
213+
? Collections.emptySet()
214+
: Collections.unmodifiableSet(systemParameters);
215+
}
216+
190217
/** Builder for SubAgentConfig. */
191218
public static class Builder {
192219
private String toolName;
@@ -196,6 +223,7 @@ public static class Builder {
196223
private Session session;
197224
private Map<String, Map<String, Object>> customParameters = new HashMap<>();
198225
private List<String> requiredCustomParameters = new ArrayList<>();
226+
private Set<String> systemParameters = new LinkedHashSet<>();
199227

200228
private Builder() {}
201229

@@ -321,6 +349,35 @@ public Builder addParameter(String name, Map<String, Object> schema, boolean req
321349
return this;
322350
}
323351

352+
/**
353+
* Adds a system parameter (e.g., userId) to be injected securely.
354+
*
355+
* <p>Unlike custom parameters, system parameters are completely invisible to the
356+
* language model and will NOT be included in the generated JSON schema. They are
357+
* extracted strictly from the {@link io.agentscope.core.tool.ToolExecutionContext}
358+
* at runtime, preventing prompt injection attacks and LLM hallucination.
359+
*
360+
* @param name The name of the system parameter
361+
* @return This builder
362+
* @throws IllegalArgumentException If the {@code name} is null, empty, or a reserved
363+
* system parameter (e.g., "message" or "session_id").
364+
*/
365+
public Builder addSystemParameter(String name) {
366+
if ("message".equals(name) || "session_id".equals(name)) {
367+
throw new IllegalArgumentException(
368+
"Cannot use reserved parameter name: '"
369+
+ name
370+
+ "'. This is a built-in parameter.");
371+
}
372+
if (name == null || name.trim().isEmpty()) {
373+
throw new IllegalArgumentException(
374+
"System parameter name cannot be null or empty.");
375+
}
376+
377+
this.systemParameters.add(name);
378+
return this;
379+
}
380+
324381
/**
325382
* Builds the configuration.
326383
*

‎agentscope-core/src/main/java/io/agentscope/core/tool/subagent/SubAgentTool.java‎

Lines changed: 69 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,15 @@
5959
* <li>{@code message} - Required. The message to send to the agent.
6060
* </ul>
6161
*
62-
* <p>Additionally, custom parameters defined in {@link SubAgentConfig} are exposed in the
63-
* tool schema. At runtime, these parameters are extracted (prioritizing {@link ToolExecutionContext}
64-
* over LLM input) and securely injected into the sub-agent's message {@code metadata}.
62+
* <p>Additionally, the tool supports two categories of external parameters defined in
63+
* {@link SubAgentConfig}:
64+
* <ul>
65+
* <li><b>System Parameters:</b> Extracted strictly from {@link ToolExecutionContext}. These are
66+
* invisible to the LLM and completely tamper-proof.
67+
* <li><b>Custom Parameters:</b> Exposed in the tool schema for LLM inference. At runtime, these
68+
* are extracted with {@link ToolExecutionContext} taking precedence over LLM input.
69+
* </ul>
70+
* Both types of parameters are securely injected into the sub-agent's message {@code metadata}.
6571
*/
6672
public class SubAgentTool implements AgentTool {
6773

@@ -124,7 +130,8 @@ public Mono<ToolResultBlock> callAsync(ToolCallParam param) {
124130
* <ul>
125131
* <li>Session ID generation for new conversations
126132
* <li>Agent state loading for continued sessions
127-
* <li>Extraction of custom parameters from ToolExecutionContext and LLM input
133+
* <li>Extraction of external parameters with strict priority:
134+
* {@link ToolExecutionContext} (highest) &gt; LLM input (lowest)
128135
* <li>Secure injection of custom parameters into the message metadata
129136
* <li>Message execution (streaming or non-streaming based on config)
130137
* <li>Agent state persistence after execution
@@ -152,47 +159,41 @@ private Mono<ToolResultBlock> executeConversation(ToolCallParam param) {
152159
return Mono.just(ToolResultBlock.error("Message is required"));
153160
}
154161

155-
// Extract custom external parameters (filter out system built-in keys)
162+
// Extract external parameters
156163
Map<String, Object> finalMetadata = new HashMap<>();
164+
ToolExecutionContext context = param.getContext();
165+
166+
// 1. Extract system parameters (Securely injected, invisible to LLM)
167+
Set<String> declaredSystemParams = config.getSystemParameters();
168+
if (declaredSystemParams != null && !declaredSystemParams.isEmpty()) {
169+
for (String paramName : declaredSystemParams) {
170+
Object ctxValue = extractFromContext(context, paramName);
171+
if (ctxValue != null) {
172+
finalMetadata.put(paramName, ctxValue);
173+
} else {
174+
logger.warn(
175+
"System parameter '{}' was declared but not provided in"
176+
+ " ToolExecutionContext",
177+
paramName);
178+
}
179+
}
180+
}
181+
182+
// 2. Extract custom parameters (LLM inferred, Context overrides LLM)
157183
Map<String, Map<String, Object>> declaredCustomParams =
158184
config.getCustomParameters();
159-
160185
if (declaredCustomParams != null && !declaredCustomParams.isEmpty()) {
161186
for (String paramName : declaredCustomParams.keySet()) {
162187
boolean injectedFromContext = false;
163188

164189
// Prioritize obtaining from ToolExecutionContext
165-
ToolExecutionContext context = param.getContext();
166-
if (context != null) {
167-
// sequentially try the standard JSON types that could be passed
168-
// to the LLM
169-
Object ctxValue = context.get(paramName, String.class);
170-
if (ctxValue == null) {
171-
ctxValue = context.get(paramName, Integer.class);
172-
}
173-
if (ctxValue == null) {
174-
ctxValue = context.get(paramName, Boolean.class);
175-
}
176-
if (ctxValue == null) {
177-
ctxValue = context.get(paramName, Double.class);
178-
}
179-
if (ctxValue == null) {
180-
ctxValue = context.get(paramName, Map.class);
181-
}
182-
if (ctxValue == null) {
183-
ctxValue = context.get(paramName, List.class);
184-
}
185-
if (ctxValue == null) {
186-
ctxValue = context.get(paramName, Object.class);
187-
}
188-
189-
if (ctxValue != null) {
190-
finalMetadata.put(paramName, ctxValue);
191-
injectedFromContext = true;
192-
}
190+
Object ctxValue = extractFromContext(context, paramName);
191+
if (ctxValue != null) {
192+
finalMetadata.put(paramName, ctxValue);
193+
injectedFromContext = true;
193194
}
194195

195-
// Attempt to retrieve from input
196+
// Attempt to retrieve from LLM input if not injected by context
196197
if (!injectedFromContext && input.containsKey(paramName)) {
197198
finalMetadata.put(paramName, input.get(paramName));
198199
}
@@ -248,6 +249,39 @@ private Mono<ToolResultBlock> executeConversation(ToolCallParam param) {
248249
});
249250
}
250251

252+
/**
253+
* Extracts a value from the ToolExecutionContext by trying standard JSON types.
254+
*
255+
* @param context The execution context
256+
* @param paramName The name of the parameter to extract
257+
* @return The extracted value, or null if not found
258+
*/
259+
private Object extractFromContext(ToolExecutionContext context, String paramName) {
260+
if (context == null) {
261+
return null;
262+
}
263+
Object ctxValue = context.get(paramName, String.class);
264+
if (ctxValue == null) {
265+
ctxValue = context.get(paramName, Integer.class);
266+
}
267+
if (ctxValue == null) {
268+
ctxValue = context.get(paramName, Boolean.class);
269+
}
270+
if (ctxValue == null) {
271+
ctxValue = context.get(paramName, Double.class);
272+
}
273+
if (ctxValue == null) {
274+
ctxValue = context.get(paramName, Map.class);
275+
}
276+
if (ctxValue == null) {
277+
ctxValue = context.get(paramName, List.class);
278+
}
279+
if (ctxValue == null) {
280+
ctxValue = context.get(paramName, Object.class);
281+
}
282+
return ctxValue;
283+
}
284+
251285
/**
252286
* Loads agent state from the session storage.
253287
*

‎agentscope-core/src/test/java/io/agentscope/core/tool/subagent/SubAgentConfigTest.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,32 @@ void testSchemaCustomParameter() {
234234
assertFalse(config.getRequiredCustomParameters().contains("role"));
235235
}
236236

237+
@Test
238+
@DisplayName("Should throw exception when adding reserved system parameter name")
239+
void testAddReservedSystemParameterThrowsException() {
240+
SubAgentConfig.Builder builder = SubAgentConfig.builder();
241+
242+
assertThrows(
243+
IllegalArgumentException.class,
244+
() -> builder.addSystemParameter("message"),
245+
"Should throw IllegalArgumentException for reserved 'message'");
246+
247+
assertThrows(
248+
IllegalArgumentException.class,
249+
() -> builder.addSystemParameter("session_id"),
250+
"Should throw IllegalArgumentException for reserved 'session_id'");
251+
}
252+
253+
@Test
254+
@DisplayName("Should throw exception when system parameter name is empty")
255+
void testAddEmptySystemParameterNameThrowsException() {
256+
SubAgentConfig.Builder builder = SubAgentConfig.builder();
257+
258+
assertThrows(IllegalArgumentException.class, () -> builder.addSystemParameter(""));
259+
260+
assertThrows(IllegalArgumentException.class, () -> builder.addSystemParameter(null));
261+
}
262+
237263
@Test
238264
@DisplayName("Builder should be chainable")
239265
void testBuilderChaining() {
@@ -246,6 +272,8 @@ void testBuilderChaining() {
246272
.forwardEvents(true)
247273
.streamOptions(null)
248274
.session(new InMemorySession())
275+
.addParameter("param", null, true)
276+
.addSystemParameter("sysParam")
249277
.build();
250278

251279
assertNotNull(config);
@@ -302,6 +330,14 @@ void testGetRequiredCustomParametersEmpty() {
302330
assertNotNull(config.getRequiredCustomParameters());
303331
assertTrue(config.getRequiredCustomParameters().isEmpty());
304332
}
333+
334+
@Test
335+
@DisplayName("getSystemParameters() should return empty set when not set")
336+
void testGetSystemParametersEmpty() {
337+
SubAgentConfig config = SubAgentConfig.builder().build();
338+
assertNotNull(config.getSystemParameters());
339+
assertTrue(config.getSystemParameters().isEmpty());
340+
}
305341
}
306342

307343
@Nested

0 commit comments

Comments
 (0)