Skip to content

fix(core): isolate active tool groups per session #2444

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
logicwu0 wants to merge 1 commit into
base: main
Choose a base branch
from
+118 −30
Open

fix(core): isolate active tool groups per session #2444

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

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 38 additions & 30 deletions agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -423,14 +423,14 @@ private static AgentState freshState(

/**
* Persist the current {@link AgentState} via the configured {@link AgentStateStore}, or {@code
* Mono.empty()} when no AgentStateStore was provided. Synchronises toolkit activeGroups into the state
* before writing.
* Mono.empty()} when no AgentStateStore was provided. Synchronises the call-scoped toolkit's
* active groups into the state before writing.
*/
private Mono<Void> saveStateToSession(CallExecution scope) {
if (stateStore == null) {
return Mono.empty();
}
syncToolkitToState(scope.state);
syncToolkitToState(scope);
SlotRef ref = SlotRef.parse(scope.slotKey);
AgentState toSave = scope.state;
return Mono.<Void>fromRunnable(
Expand Down Expand Up @@ -494,9 +494,6 @@ private CallExecution activateSlotForContext(RuntimeContext ctx) {
slot, k -> new PermissionEngine(loaded.getPermissionContext()));
}
CallExecution scope = new CallExecution(loaded, loadedEngine, slot);
if (toolkit != null) {
toolkit.setActiveGroups(loaded.getToolContext().getActivatedGroups());
}
return scope;
}

Expand Down Expand Up @@ -615,7 +612,7 @@ private Mono<String> applySystemPromptMiddlewares(String prompt, RuntimeContext
protected void consumeSystemMsgAfterPreCall(Msg systemMsg, Object callScope) {
CallExecution ce = (CallExecution) callScope;
ce.systemMsg = systemMsg;
syncToolkitToState(ce.state);
syncToolkitToState(ce);
}

@Override
Expand Down Expand Up @@ -1051,25 +1048,33 @@ private Mono<Msg> doStructuredCall(List<Msg> msgs, Class<?> targetClass, JsonNod
targetClass != null
? JsonSchemaUtils.generateSchemaFromClass(targetClass)
: JsonSchemaUtils.generateSchemaFromJsonNode(schemaDesc);
boolean hasTools = !toolkit.getToolSchemas().isEmpty();
boolean useNative =
hasTools
? model.supportsNativeStructuredOutputWithTools()
: model.supportsNativeStructuredOutput();
if (useNative) {
return doNativeStructuredCall(msgs, jsonSchema)
.onErrorResume(
e -> {
log.warn(
"Native structured output failed ({}) — falling back to"
+ " synthetic tool path",
e.getMessage() != null
? e.getMessage()
: e.getClass().getSimpleName());
return doFallbackStructuredCall(msgs, jsonSchema);
});
}
return doFallbackStructuredCall(msgs, jsonSchema);
return Mono.deferContextual(
cv -> {
CallExecution scope = scopeFrom(cv);
boolean hasTools =
!scope.toolkit
.getToolSchemas(
scope.state.getToolContext().getActivatedGroups())
.isEmpty();
boolean useNative =
hasTools
? model.supportsNativeStructuredOutputWithTools()
: model.supportsNativeStructuredOutput();
if (useNative) {
return doNativeStructuredCall(msgs, jsonSchema)
.onErrorResume(
e -> {
log.warn(
"Native structured output failed ({}) — falling"
+ " back to synthetic tool path",
e.getMessage() != null
? e.getMessage()
: e.getClass().getSimpleName());
return doFallbackStructuredCall(msgs, jsonSchema);
});
}
return doFallbackStructuredCall(msgs, jsonSchema);
});
}

/**
Expand Down Expand Up @@ -1437,6 +1442,7 @@ final class CallExecution {
AgentState state;
PermissionEngine permissionEngine;
String slotKey;
final Toolkit toolkit;

/**
* Per-call system message, propagated across PreCallEvent → PreReasoningEvent /
Expand Down Expand Up @@ -1490,6 +1496,8 @@ final class CallExecution {
this.state = state;
this.permissionEngine = permissionEngine;
this.slotKey = slotKey;
this.toolkit = ReActAgent.this.toolkit.copy();
this.toolkit.setActiveGroups(state.getToolContext().getActivatedGroups());
}

/**
Expand Down Expand Up @@ -2417,7 +2425,7 @@ private Mono<Msg> acting(int iter) {
buildSuspendedMsg(pendingPairs));
}

syncToolkitToState(state);
syncToolkitToState(CallExecution.this);
return executeIteration(iter + 1);
});
});
Expand Down Expand Up @@ -3772,9 +3780,9 @@ public String getDefaultSessionId() {
return defaultSessionId;
}

private void syncToolkitToState(AgentState state) {
if (toolkit != null && state != null) {
state.getToolContext().setActivatedGroups(toolkit.getActiveGroups());
private void syncToolkitToState(CallExecution scope) {
if (scope != null && scope.state != null) {
scope.state.getToolContext().setActivatedGroups(scope.toolkit.getActiveGroups());
}
}

Expand Down
80 changes: 80 additions & 0 deletions agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,14 @@
import io.agentscope.core.message.Msg;
import io.agentscope.core.message.MsgRole;
import io.agentscope.core.message.TextBlock;
import io.agentscope.core.middleware.MiddlewareBase;
import io.agentscope.core.model.ChatModelBase;
import io.agentscope.core.model.ChatResponse;
import io.agentscope.core.model.GenerateOptions;
import io.agentscope.core.model.ToolSchema;
import io.agentscope.core.state.AgentState;
import io.agentscope.core.state.InMemoryAgentStateStore;
import io.agentscope.core.state.ToolContextState;
import io.agentscope.core.state.legacy.ToolkitState;
import io.agentscope.core.tool.Toolkit;
import java.time.Duration;
Expand Down Expand Up @@ -316,6 +318,84 @@ void concurrentDistinctSessionsAreIsolated() {
}
}

@Test
@DisplayName("concurrent sessions preserve their own active tool groups")
void concurrentSessionsPreserveActiveToolGroups() {
InMemoryAgentStateStore store = new InMemoryAgentStateStore();
store.save(
"u",
"session-a",
"agent_state",
AgentState.builder()
.userId("u")
.sessionId("session-a")
.toolContext(
ToolContextState.builder().addActivatedGroup("group-a").build())
.build());
store.save(
"u",
"session-b",
"agent_state",
AgentState.builder()
.userId("u")
.sessionId("session-b")
.toolContext(
ToolContextState.builder().addActivatedGroup("group-b").build())
.build());

Toolkit toolkit = new Toolkit();
toolkit.createToolGroup("group-a", "Session A tools", false);
toolkit.createToolGroup("group-b", "Session B tools", false);
CountDownLatch bothCallsActivated = new CountDownLatch(2);
MiddlewareBase activationBarrier =
new MiddlewareBase() {
@Override
public Mono<String> onSystemPrompt(
Agent agent, RuntimeContext ctx, String currentPrompt) {
return Mono.fromCallable(
() -> {
bothCallsActivated.countDown();
assertTrue(
bothCallsActivated.await(5, TimeUnit.SECONDS),
"both calls should activate their session state");
return currentPrompt;
});
}
};
ReActAgent agent =
ReActAgent.builder()
.name("asst")
.sysPrompt("hi")
.model(new NoopModel())
.toolkit(toolkit)
.stateStore(store)
.middleware(activationBarrier)
.build();

Mono<Msg> callA =
agent.call(
List.of(userMsg("hello-a")),
RuntimeContext.builder().userId("u").sessionId("session-a").build())
.subscribeOn(Schedulers.parallel());
Mono<Msg> callB =
agent.call(
List.of(userMsg("hello-b")),
RuntimeContext.builder().userId("u").sessionId("session-b").build())
.subscribeOn(Schedulers.parallel());

Mono.when(callA, callB).block(Duration.ofSeconds(10));

assertEquals(
List.of("group-a"),
agent.getAgentState("u", "session-a").getToolContext().getActivatedGroups());
assertEquals(
List.of("group-b"),
agent.getAgentState("u", "session-b").getToolContext().getActivatedGroups());
assertTrue(
agent.getToolkit().getActiveGroups().isEmpty(),
"session activation must not mutate the shared toolkit");
}

@Test
@DisplayName("concurrent calls to the same session are serialized (no lost updates)")
void concurrentSameSessionIsSerialized() {
Expand Down
Loading