diff --git a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java index 003bbfeb15..47cb932007 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java +++ b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java @@ -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 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.fromRunnable( @@ -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; } @@ -615,7 +612,7 @@ private Mono applySystemPromptMiddlewares(String prompt, RuntimeContext protected void consumeSystemMsgAfterPreCall(Msg systemMsg, Object callScope) { CallExecution ce = (CallExecution) callScope; ce.systemMsg = systemMsg; - syncToolkitToState(ce.state); + syncToolkitToState(ce); } @Override @@ -1051,25 +1048,33 @@ private Mono doStructuredCall(List 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); + }); } /** @@ -1437,6 +1442,7 @@ final class CallExecution { AgentState state; PermissionEngine permissionEngine; String slotKey; + final Toolkit toolkit; /** * Per-call system message, propagated across PreCallEvent → PreReasoningEvent / @@ -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()); } /** @@ -2417,7 +2425,7 @@ private Mono acting(int iter) { buildSuspendedMsg(pendingPairs)); } - syncToolkitToState(state); + syncToolkitToState(CallExecution.this); return executeIteration(iter + 1); }); }); @@ -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()); } } diff --git a/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java b/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java index 4f7a8a09e5..9a78109d36 100644 --- a/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java +++ b/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java @@ -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; @@ -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 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 callA = + agent.call( + List.of(userMsg("hello-a")), + RuntimeContext.builder().userId("u").sessionId("session-a").build()) + .subscribeOn(Schedulers.parallel()); + Mono 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() {