Skip to content
76 changes: 56 additions & 20 deletions agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@
import io.agentscope.core.model.ToolSchema;
import io.agentscope.core.permission.PermissionBehavior;
import io.agentscope.core.permission.PermissionContextState;
import io.agentscope.core.permission.PermissionDecision;
import io.agentscope.core.permission.PermissionEngine;
import io.agentscope.core.permission.PermissionMode;
import io.agentscope.core.permission.PermissionRule;
Expand Down Expand Up @@ -3070,6 +3071,14 @@ Flux<AgentEvent> actingStream(
.flatMapMany(
gate -> {
List<ToolUseBlock> pending = gate.pendingAsk();
Map<String, List<PermissionRule>> pendToolMap =
pending.stream()
.collect(
Collectors.toMap(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Collectors.toMap throws NPE when the value mapper returns null. Today ToolUseBlock#getSuggestedRules normalises null to an empty list so this is safe, but it is only safe because of a guarantee one class away — an explicit loop (or Objects.requireNonNullElseGet) keeps this from breaking if a legacy/deserialized block ever yields null.

ToolUseBlock::getId,
ToolUseBlock::getSuggestedRules,
(a, b) -> a));

Set<String> autoDenied = gate.autoDeniedIds();

// Mark ToolUseBlock.state in context for every gated tool. ALLOWED
Expand All @@ -3084,17 +3093,11 @@ Flux<AgentEvent> actingStream(
}
stateUpdates.put(
tc.getId(),
pending.stream()
.anyMatch(
p ->
p.getId()
.equals(
tc
.getId()))
pendToolMap.containsKey(tc.getId())
? ToolCallState.ASKING
: ToolCallState.ALLOWED);
}
updateToolCallStates(stateUpdates);
updateToolCallStatesAndRules(stateUpdates, pendToolMap);

if (pending.isEmpty()) {
return runToolBatch(
Expand Down Expand Up @@ -3431,31 +3434,55 @@ private Mono<PermissionVerdict> evaluateOne(ToolUseBlock use, boolean useEngine)
return permissionEngine
.checkPermission(tb, input)
.map(
decision ->
new PermissionVerdict(
use,
decision == null
? PermissionBehavior.ASK
: decision.getBehavior()));
decision -> {
if (decision == null) {
return new PermissionVerdict(use, PermissionBehavior.ASK);
}

// Carry the engine's suggested rules on the ToolUseBlock
ToolUseBlock toolUseBlock =
withSuggestedRulesIfExist(use, decision);
return new PermissionVerdict(
toolUseBlock, decision.getBehavior());
});
}
return tb.checkPermissions(input, state.getPermissionContext())
.map(
decision -> {
if (decision == null) {
return new PermissionVerdict(use, PermissionBehavior.ALLOW);
}

// Carry the engine's suggested rules on the ToolUseBlock
ToolUseBlock newToolUse = withSuggestedRulesIfExist(use, decision);

// In the legacy lightweight path only an explicit ASK from the tool
// gates execution; PASSTHROUGH and ALLOW both run, DENY is
// honoured.
return switch (decision.getBehavior()) {
case ASK -> new PermissionVerdict(use, PermissionBehavior.ASK);
case ASK ->
new PermissionVerdict(
newToolUse, PermissionBehavior.ASK);
case DENY ->
new PermissionVerdict(use, PermissionBehavior.DENY);
default -> new PermissionVerdict(use, PermissionBehavior.ALLOW);
new PermissionVerdict(
newToolUse, PermissionBehavior.DENY);
default ->
new PermissionVerdict(
newToolUse, PermissionBehavior.ALLOW);
};
});
}

private ToolUseBlock withSuggestedRulesIfExist(
ToolUseBlock use, PermissionDecision decision) {
// Carry the engine's suggested rules on the ToolUseBlock
List<PermissionRule> suggested = decision.getSuggestedRules();
if (suggested == null || suggested.isEmpty()) {
return use;
}
return use.withSuggestedRules(suggested);
}

private record PermissionVerdict(ToolUseBlock use, PermissionBehavior behavior) {}

private List<ToolUseBlock> getSuspendedToolCalls(
Expand Down Expand Up @@ -4281,11 +4308,12 @@ private List<ToolUseBlock> extractRecentToolCalls() {
// ==================== Tool call state helpers (Permission HITL) ====================

/**
* Locate the last assistant Msg in context and replace the {@code state} of every
* Locate the last assistant Msg in context and replace the {@code state,suggestedRules} of every
* {@link ToolUseBlock} whose id matches the given map's key. Mirrors Python's
* {@code _update_tool_call_state} but operates in bulk to minimise list rebuilds.
*/
private void updateToolCallStates(Map<String, ToolCallState> updates) {
private void updateToolCallStatesAndRules(
Map<String, ToolCallState> updates, Map<String, List<PermissionRule>> pendToolMap) {
if (updates == null || updates.isEmpty()) {
return;
}
Expand All @@ -4297,12 +4325,20 @@ private void updateToolCallStates(Map<String, ToolCallState> updates) {
for (ToolUseBlock toolUse : lastAssistant.getContentBlocks(ToolUseBlock.class)) {
ToolCallState newState = updates.get(toolUse.getId());
if (newState != null) {
replacements.put(toolUse.getId(), toolUse.withState(newState));
List<PermissionRule> rules =
pendToolMap.getOrDefault(toolUse.getId(), List.of());
replacements.put(
toolUse.getId(), toolUse.withStateAndSuggestedRules(newState, rules));
}
}
MessageUtils.replaceToolUseBlocks(state.contextMutable(), replacements);
}

/** Convenience overload for a single tool call. */
private void updateToolCallState(String toolCallId, ToolCallState newState) {
updateToolCallStatesAndRules(Map.of(toolCallId, newState), Map.of());
}

/** The ToolUseBlocks in the last assistant Msg that are in ASKING state (HITL pending). */
private List<ToolUseBlock> askingToolCalls() {
Msg last = MessageUtils.lastAssistantMessage(state.contextMutable());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,11 @@
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
import io.agentscope.core.permission.PermissionRule;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
Expand Down Expand Up @@ -51,6 +54,7 @@ public final class ToolUseBlock extends ContentBlock {
private final String content; // Raw content for streaming tool calls
private final Map<String, Object> metadata; // Provider-specific metadata
private final ToolCallState state;
private final List<PermissionRule> suggestedRules; // Rules suggested when gated by ASK

/**
* Creates a new tool use block with metadata (convenience constructor).
Expand All @@ -62,7 +66,7 @@ public final class ToolUseBlock extends ContentBlock {
*/
public ToolUseBlock(
String id, String name, Map<String, Object> input, Map<String, Object> metadata) {
this(id, name, input, null, metadata, null);
this(id, name, input, null, metadata, null, null);
}

/**
Expand All @@ -73,7 +77,7 @@ public ToolUseBlock(
* @param input Input parameters for the tool (will be defensively copied)
*/
public ToolUseBlock(String id, String name, Map<String, Object> input) {
this(id, name, input, null, null, null);
this(id, name, input, null, null, null, null);
}

/**
Expand All @@ -91,7 +95,7 @@ public ToolUseBlock(
Map<String, Object> input,
String content,
Map<String, Object> metadata) {
this(id, name, input, content, metadata, null);
this(id, name, input, content, metadata, null, null);
}

/**
Expand All @@ -104,14 +108,36 @@ public ToolUseBlock(
* @param metadata Provider-specific metadata (will be defensively copied)
* @param state The tool call state, defaults to PENDING if null
*/
public ToolUseBlock(
String id,
String name,
Map<String, Object> input,
String content,
Map<String, Object> metadata,
ToolCallState state) {
this(id, name, input, content, metadata, state, null);
}

/**
* Creates a new tool use block with all fields.
*
* @param id Unique identifier for this tool call
* @param name Name of the tool to execute
* @param input Input parameters for the tool (will be defensively copied)
* @param content Raw content for streaming tool calls
* @param metadata Provider-specific metadata (will be defensively copied)
* @param state The tool call state, defaults to PENDING if null
* @param suggestedRules Permission rules suggested for this call (will be defensively copied)
*/
@JsonCreator
public ToolUseBlock(
@JsonProperty("id") String id,
@JsonProperty("name") String name,
@JsonProperty("input") Map<String, Object> input,
@JsonProperty("content") String content,
@JsonProperty("metadata") Map<String, Object> metadata,
@JsonProperty("state") ToolCallState state) {
@JsonProperty("state") ToolCallState state,
@JsonProperty("suggested_rules") List<PermissionRule> suggestedRules) {
this.id = id;
this.name = name;
// Defensive copy to prevent external modifications
Expand All @@ -125,6 +151,10 @@ public ToolUseBlock(
? Collections.emptyMap()
: Collections.unmodifiableMap(new HashMap<>(metadata));
this.state = state != null ? state : ToolCallState.PENDING;
this.suggestedRules =
suggestedRules == null
? Collections.emptyList()
: Collections.unmodifiableList(new ArrayList<>(suggestedRules));
}

/**
Expand Down Expand Up @@ -184,6 +214,23 @@ public ToolCallState getState() {
return state;
}

/**
* Gets the permission rules suggested for this tool call.
*
* <p>Populated when the permission engine gates this call with
* {@link io.agentscope.core.permission.PermissionBehavior#ASK}: the owning tool derives
* candidate rules from the actual invocation via
* {@link io.agentscope.core.tool.ToolBase#generateSuggestions}. Passing them back in a
* {@code ConfirmResult} registers them with the engine, so future matching calls resolve
* without another prompt.
*
* @return The suggested rules, or an empty list when none were suggested
*/
@JsonProperty("suggested_rules")
public List<PermissionRule> getSuggestedRules() {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two small things: (1) the name pendToolMap reads as key→tool while it is key→rules — pendingSuggestedRules would match the new getter; (2) getSuggestedRules() documents "empty list when none were suggested", so the caller-side assertNotNull in ReActAgentHitlTest tests the contract of a method that cannot return null — asserting the context block (not only the RequireUserConfirmEvent payload) carries the rules would cover the write-back path in updateToolCallStatesAndRules, which is where the [Warning] above lives.

return suggestedRules;
}

/**
* Checks whether this tool call is executed by the provider on the server side (marked via
* {@link #METADATA_SERVER_TOOL}).
Expand All @@ -202,7 +249,43 @@ public boolean isServerTool() {
* @return A new ToolUseBlock with the updated state
*/
public ToolUseBlock withState(ToolCallState state) {
return new ToolUseBlock(this.id, this.name, this.input, this.content, this.metadata, state);
return new ToolUseBlock(
this.id,
this.name,
this.input,
this.content,
this.metadata,
state,
this.suggestedRules);
}

/**
* Returns a copy of this block with the given suggested rules.
*
* @param suggestedRules The suggested permission rules
* @return A new ToolUseBlock with the updated suggested rules
*/
public ToolUseBlock withSuggestedRules(List<PermissionRule> suggestedRules) {
return new ToolUseBlock(
this.id,
this.name,
this.input,
this.content,
this.metadata,
this.state,
suggestedRules);
}

/**
* Returns a copy of this block with the given suggested rules.
* @param state The new state
* @param suggestedRules The suggested permission rules
* @return A new ToolUseBlock with the updated suggested rules and state
*/
public ToolUseBlock withStateAndSuggestedRules(
ToolCallState state, List<PermissionRule> suggestedRules) {
return new ToolUseBlock(
this.id, this.name, this.input, this.content, this.metadata, state, suggestedRules);
}

/**
Expand All @@ -224,6 +307,7 @@ public static class Builder {
private String content;
private Map<String, Object> metadata;
private ToolCallState state;
private List<PermissionRule> suggestedRules;

/**
* Sets the unique identifier for the tool call.
Expand Down Expand Up @@ -294,13 +378,24 @@ public Builder state(ToolCallState state) {
return this;
}

/**
* Sets the permission rules suggested for this tool call.
*
* @param suggestedRules The suggested permission rules
* @return This builder for chaining
*/
public Builder suggestedRules(List<PermissionRule> suggestedRules) {
this.suggestedRules = suggestedRules;
return this;
}

/**
* Builds a new ToolUseBlock with the configured properties.
*
* @return A new ToolUseBlock instance
*/
public ToolUseBlock build() {
return new ToolUseBlock(id, name, input, content, metadata, state);
return new ToolUseBlock(id, name, input, content, metadata, state, suggestedRules);
}
}
}
Loading
Loading