Skip to content

Commit 44056ef

Browse files
fang-techAlbumenJ
andauthored
feat(tool): add thread-safe whitelist management and dynamic desc (#379)
## AgentScope-Java Version 1.0.4 ## Description Add a list of allowed commands to the description sent to LLM. This helps to reduce the occurrence of LLM running commands that are not allowed. Although it incurs an increase in the consumption of one-time tokens, it can reduce the number of tool invocations, which is a more token-consuming behavior. At the same time, make the modification of the allowedCommand set thread-safe. ## Checklist Please check the following items before code is ready to be reviewed. - [x] Code has been formatted with `mvn spotless:apply` - [x] All tests are passing (`mvn test`) - [x] Javadoc comments are complete and follow project conventions - [x] Related documentation has been updated (e.g. links, examples, etc.) - [x] Code is ready for review --------- Co-authored-by: Albumen Kevin <jhq0812@gmail.com>
1 parent cacf502 commit 44056ef

2 files changed

Lines changed: 421 additions & 32 deletions

File tree

agentscope-core/src/main/java/io/agentscope/core/tool/coding/ShellCommandTool.java

Lines changed: 143 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -17,18 +17,23 @@
1717

1818
import io.agentscope.core.message.TextBlock;
1919
import io.agentscope.core.message.ToolResultBlock;
20-
import io.agentscope.core.tool.Tool;
21-
import io.agentscope.core.tool.ToolParam;
20+
import io.agentscope.core.tool.AgentTool;
21+
import io.agentscope.core.tool.ToolCallParam;
2222
import java.io.BufferedReader;
2323
import java.io.IOException;
2424
import java.io.InputStreamReader;
2525
import java.nio.charset.StandardCharsets;
2626
import java.time.Duration;
27-
import java.util.HashSet;
27+
import java.util.ArrayList;
28+
import java.util.Collections;
29+
import java.util.List;
30+
import java.util.Map;
2831
import java.util.Set;
32+
import java.util.concurrent.ConcurrentHashMap;
2933
import java.util.concurrent.TimeUnit;
3034
import java.util.concurrent.TimeoutException;
3135
import java.util.function.Function;
36+
import java.util.stream.Collectors;
3237
import org.slf4j.Logger;
3338
import org.slf4j.LoggerFactory;
3439
import reactor.core.publisher.Mono;
@@ -48,7 +53,7 @@
4853
* @see UnixCommandValidator
4954
* @see WindowsCommandValidator
5055
*/
51-
public class ShellCommandTool {
56+
public class ShellCommandTool implements AgentTool {
5257

5358
private static final Logger logger = LoggerFactory.getLogger(ShellCommandTool.class);
5459
private static final int DEFAULT_TIMEOUT = 300;
@@ -87,9 +92,14 @@ public ShellCommandTool(
8792
Set<String> allowedCommands,
8893
Function<String, Boolean> approvalCallback,
8994
CommandValidator commandValidator) {
90-
// If allowedCommands is null, create an empty HashSet (which means allow all by default)
91-
// If provided, use it directly
92-
this.allowedCommands = allowedCommands != null ? allowedCommands : new HashSet<>();
95+
// Use ConcurrentHashMap.newKeySet() for thread-safe, high-performance concurrent access
96+
// Create defensive copy to prevent external modifications
97+
if (allowedCommands != null && !allowedCommands.isEmpty()) {
98+
this.allowedCommands = ConcurrentHashMap.newKeySet(allowedCommands.size());
99+
this.allowedCommands.addAll(allowedCommands);
100+
} else {
101+
this.allowedCommands = ConcurrentHashMap.newKeySet();
102+
}
93103
this.approvalCallback = approvalCallback;
94104
this.commandValidator =
95105
commandValidator != null ? commandValidator : createDefaultValidator();
@@ -109,16 +119,137 @@ private static CommandValidator createDefaultValidator() {
109119
}
110120
}
111121

122+
// =============================== Allowed Commands Management ===============================
123+
112124
/**
113-
* Get the set of allowed commands.
114-
* The returned set can be modified to dynamically update the whitelist.
125+
* Get an unmodifiable view of the allowed commands.
126+
* The returned set is thread-safe for reading but cannot be modified directly.
127+
* Use {@link #addAllowedCommand(String)}, {@link #removeAllowedCommand(String)},
128+
* or {@link #clearAllowedCommands()} to modify the whitelist.
115129
*
116-
* @return The mutable set of allowed command executables
130+
* @return An unmodifiable view of the allowed command executables
117131
*/
118132
public Set<String> getAllowedCommands() {
119-
return allowedCommands;
133+
return Collections.unmodifiableSet(allowedCommands);
134+
}
135+
136+
/**
137+
* Add a command to the whitelist in a thread-safe manner.
138+
*
139+
* @param command The command executable to add
140+
* @return true if the command was added, false if it was already present
141+
*/
142+
public boolean addAllowedCommand(String command) {
143+
if (command == null || command.trim().isEmpty()) {
144+
throw new IllegalArgumentException("Command cannot be null or empty");
145+
}
146+
boolean added = allowedCommands.add(command);
147+
if (added) {
148+
logger.debug("Added command to whitelist: {}", command);
149+
}
150+
return added;
151+
}
152+
153+
/**
154+
* Remove a command from the whitelist in a thread-safe manner.
155+
*
156+
* @param command The command executable to remove
157+
* @return true if the command was removed, false if it was not present
158+
*/
159+
public boolean removeAllowedCommand(String command) {
160+
boolean removed = allowedCommands.remove(command);
161+
if (removed) {
162+
logger.debug("Removed command from whitelist: {}", command);
163+
}
164+
return removed;
165+
}
166+
167+
/**
168+
* Clear all commands from the whitelist in a thread-safe manner.
169+
*/
170+
public void clearAllowedCommands() {
171+
allowedCommands.clear();
172+
logger.debug("Cleared all commands from whitelist");
120173
}
121174

175+
/**
176+
* Check if a command is in the whitelist.
177+
*
178+
* @param command The command executable to check
179+
* @return true if the command is whitelisted
180+
*/
181+
public boolean isCommandAllowed(String command) {
182+
return allowedCommands.contains(command);
183+
}
184+
185+
// ========================= AgentTool interface implementation =========================
186+
187+
@Override
188+
public String getName() {
189+
return "execute_shell_command";
190+
}
191+
192+
@Override
193+
public String getDescription() {
194+
StringBuilder desc = new StringBuilder();
195+
desc.append("Execute a shell command with security validation and return the result.");
196+
197+
// Add whitelist information if configured
198+
if (!allowedCommands.isEmpty()) {
199+
desc.append(" ALLOWED COMMANDS WHITELIST: [");
200+
String commandList =
201+
new ArrayList<>(allowedCommands).stream().collect(Collectors.joining(", "));
202+
desc.append(commandList);
203+
desc.append("]. Only these commands can be executed directly.");
204+
} else {
205+
desc.append(" No whitelist configured - all commands require approval.");
206+
}
207+
208+
desc.append(" Commands are validated against the whitelist (if configured).");
209+
desc.append(" Non-whitelisted commands require user approval via callback.");
210+
desc.append(
211+
" Multiple command separators (&, |, ;) are detected and blocked for security.");
212+
desc.append(" Returns output in format:");
213+
desc.append(" <returncode>code</returncode><stdout>output</stdout><stderr>error</stderr>.");
214+
desc.append(" If command is rejected, returncode will be -1 with SecurityError in stderr.");
215+
216+
return desc.toString();
217+
}
218+
219+
@Override
220+
public Map<String, Object> getParameters() {
221+
return Map.of(
222+
"type", "object",
223+
"properties",
224+
Map.of(
225+
"command",
226+
Map.of(
227+
"type",
228+
"string",
229+
"description",
230+
"The shell command to execute"),
231+
"timeout",
232+
Map.of(
233+
"type",
234+
"integer",
235+
"description",
236+
"The maximum time (in seconds) allowed for the"
237+
+ " command to run (default: 300)")),
238+
"required", List.of("command"));
239+
}
240+
241+
@Override
242+
public Mono<ToolResultBlock> callAsync(ToolCallParam param) {
243+
Map<String, Object> input = param.getInput();
244+
String command = (String) input.get("command");
245+
Integer timeout =
246+
input.containsKey("timeout") ? ((Number) input.get("timeout")).intValue() : null;
247+
248+
return executeShellCommand(command, timeout);
249+
}
250+
251+
// =============================== Execute shell command ===============================
252+
122253
/**
123254
* Execute a shell command and return the return code, standard output, and
124255
* standard error within tags.
@@ -135,26 +266,7 @@ public Set<String> getAllowedCommands() {
135266
* @param timeout The maximum time (in seconds) allowed for the command to run (default: 300)
136267
* @return A ToolResultBlock containing the formatted output with returncode, stdout, and stderr
137268
*/
138-
@Tool(
139-
name = "execute_shell_command",
140-
description =
141-
"Execute a shell command with security validation and return the result."
142-
+ " Commands are validated against a whitelist (if configured)."
143-
+ " Non-whitelisted commands require user approval via callback. Multiple"
144-
+ " command separators (&, |, ;) are detected and blocked for security."
145-
+ " Returns output in format:"
146-
+ " <returncode>code</returncode><stdout>output</stdout><stderr>error</stderr>."
147-
+ " If command is rejected, returncode will be -1 with SecurityError in"
148-
+ " stderr.")
149-
public Mono<ToolResultBlock> executeShellCommand(
150-
@ToolParam(name = "command", description = "The shell command to execute")
151-
String command,
152-
@ToolParam(
153-
name = "timeout",
154-
description =
155-
"The maximum time (in seconds) allowed for the command to run",
156-
required = false)
157-
Integer timeout) {
269+
public Mono<ToolResultBlock> executeShellCommand(String command, Integer timeout) {
158270

159271
int actualTimeout = timeout != null && timeout > 0 ? timeout : DEFAULT_TIMEOUT;
160272
logger.debug(

0 commit comments

Comments
 (0)