fix(generate_response): Fix flaky concurrent tests for generate_response

agentscope-core/src/main/java/io/agentscope/core/agent/StructuredOutputCapableAgent.java

@@ -189,9 +189,9 @@ private Mono<Msg> executeWithStructuredOutput(
                                     })
                             .doFinally(
                                     signal -> {
-                                        // Cleanup: remove hook and unregister tool
                                         removeHook(hook);
-                                        toolkit.removeTool(STRUCTURED_OUTPUT_TOOL_NAME);
+                                        toolkit.removeToolIfSame(
+                                                STRUCTURED_OUTPUT_TOOL_NAME, structuredOutputTool);
                                     });
                 });
     }

agentscope-core/src/main/java/io/agentscope/core/tool/ToolRegistry.java

@@ -102,6 +102,22 @@ void removeTool(String toolName) {
         registeredTools.remove(toolName);
     }
 
+    /**
+     * Atomically remove a tool only if the current instance matches the expected one.
+     * Uses {@link ConcurrentHashMap#remove(Object, Object)} to avoid TOCTOU races.
+     *
+     * @param toolName Tool name to remove
+     * @param expected The expected AgentTool instance (identity comparison)
+     * @return true if the tool was removed, false if it was already replaced or absent
+     */
+    boolean removeToolIfSame(String toolName, AgentTool expected) {
+        boolean removed = tools.remove(toolName, expected);
+        if (removed) {
+            registeredTools.remove(toolName);
+        }
+        return removed;
+    }
+
     /**
      * Remove multiple tools by names.
      *

agentscope-core/src/main/java/io/agentscope/core/tool/Toolkit.java

@@ -601,6 +601,21 @@ public void removeTool(String toolName) {
         toolRegistry.removeTool(toolName);
     }
 
+    /**
+     * Atomically remove a tool only if the registered instance is the expected one.
+     *
+     * @param toolName Name of the tool to remove
+     * @param expected The expected AgentTool instance (identity comparison)
+     * @return true if the tool was removed, false if it was already replaced or absent
+     */
+    public boolean removeToolIfSame(String toolName, AgentTool expected) {
+        if (!config.isAllowToolDeletion()) {
+            logger.warn("Tool deletion is disabled - ignoring removal of tool: {}", toolName);
+            return false;
+        }
+        return toolRegistry.removeToolIfSame(toolName, expected);
+    }
+
     /**
      * Remove tool groups and all tools within them.
      *

agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentStructuredOutputTest.java

@@ -37,7 +37,10 @@
 import java.util.List;
 import java.util.Map;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfSystemProperty;
 import reactor.core.scheduler.Schedulers;
 
 class ReActAgentStructuredOutputTest {
@@ -536,6 +539,27 @@ void testStructuredOutputPreservesThinkingBlock() {
 
     @Test
     void testConcurrencyConflictStructuredOutput() {
+        runSubscribeOnThenSequentialSecondCallStructuredOutputScenario(toolkit);
+    }
+
+    /**
+     * Reproduces the subscribeOn(elastic) vs delayed {@code doFinally} race: run many times until
+     * failure or increase confidence after a fix.
+     */
+    @EnabledIfSystemProperty(named = "agentscope.runStructuredOutputRaceTest", matches = "true")
+    @RepeatedTest(25000)
+    @DisplayName("Structured output race: 25000 repetitions (subscribeOn + immediate second call)")
+    void testConcurrencyConflictStructuredOutput_repeated() {
+        runSubscribeOnThenSequentialSecondCallStructuredOutputScenario(toolkit);
+    }
+
+    /**
+     * First call on {@link Schedulers#boundedElastic()}, second call immediately on the calling
+     * thread — same agent and toolkit. Flaky when structured-output cleanup races the second
+     * registration.
+     */
+    private void runSubscribeOnThenSequentialSecondCallStructuredOutputScenario(
+            Toolkit agentToolkit) {
         Memory memory = new InMemoryMemory();
         Map<String, Object> toolInput =
                 Map.of(
@@ -587,7 +611,7 @@ void testConcurrencyConflictStructuredOutput() {
                         .name("weather-agent")
                         .sysPrompt("You are a weather assistant")
                         .model(mockModel)
-                        .toolkit(toolkit)
+                        .toolkit(agentToolkit)
                         .memory(memory)
                         .build();
 
@@ -604,7 +628,7 @@ void testConcurrencyConflictStructuredOutput() {
         Msg responseMsg =
                 agent.call(inputMsg, WeatherResponse.class)
                         .subscribeOn(Schedulers.boundedElastic())
-                        .block(Duration.ofMillis(TestConstants.DEFAULT_TEST_TIMEOUT_MS));
+                        .block(Duration.ofMillis(TestConstants.DEFAULT_TEST_TIMEOUT_MS * 10000));
 
         Msg response2 =
                 agent.call(inputMsg, WeatherResponse.class)

agentscope-core/src/test/java/io/agentscope/core/tool/ToolkitTest.java

@@ -19,6 +19,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
@@ -241,6 +242,89 @@ void testRemoveToolWhenDeletionDisabled() {
         assertEquals(initialCount, toolkit.getToolNames().size(), "Tool count should not change");
     }
 
+    private static AgentTool namedAgentTool(String name) {
+        return new AgentTool() {
+            @Override
+            public String getName() {
+                return name;
+            }
+
+            @Override
+            public String getDescription() {
+                return "test";
+            }
+
+            @Override
+            public Map<String, Object> getParameters() {
+                return Map.of("type", "object");
+            }
+
+            @Override
+            public Mono<ToolResultBlock> callAsync(ToolCallParam param) {
+                return Mono.empty();
+            }
+        };
+    }
+
+    @Test
+    @DisplayName("removeToolIfSame removes when registered instance matches")
+    void removeToolIfSame_removesWhenInstanceMatches() {
+        AgentTool tool = namedAgentTool("remove_if_same_a");
+        toolkit.registerAgentTool(tool);
+        assertTrue(toolkit.removeToolIfSame("remove_if_same_a", tool));
+        assertEquals(null, toolkit.getTool("remove_if_same_a"));
+    }
+
+    @Test
+    @DisplayName("removeToolIfSame does not remove when a newer tool replaced the name")
+    void removeToolIfSame_noOpWhenReplaced() {
+        AgentTool first = namedAgentTool("remove_if_same_b");
+        AgentTool second = namedAgentTool("remove_if_same_b");
+        toolkit.registerAgentTool(first);
+        toolkit.registerAgentTool(second);
+        assertFalse(
+                toolkit.removeToolIfSame("remove_if_same_b", first),
+                "stale instance after replace must return false");
+        assertSame(second, toolkit.getTool("remove_if_same_b"));
+        assertTrue(toolkit.removeToolIfSame("remove_if_same_b", second));
+        assertEquals(null, toolkit.getTool("remove_if_same_b"));
+    }
+
+    @Test
+    @DisplayName("removeToolIfSame returns false when tool name is absent")
+    void removeToolIfSame_falseWhenAbsent() {
+        AgentTool phantom = namedAgentTool("remove_if_same_missing");
+        assertFalse(
+                toolkit.removeToolIfSame("remove_if_same_missing", phantom),
+                "no registration must return false");
+        assertEquals(null, toolkit.getTool("remove_if_same_missing"));
+    }
+
+    @Test
+    @DisplayName("removeToolIfSame returns false when expected is not the registered instance")
+    void removeToolIfSame_falseWhenExpectedNotRegisteredInstance() {
+        AgentTool registered = namedAgentTool("remove_if_same_wrong_ref");
+        AgentTool otherSameName = namedAgentTool("remove_if_same_wrong_ref");
+        toolkit.registerAgentTool(registered);
+        assertFalse(
+                toolkit.removeToolIfSame("remove_if_same_wrong_ref", otherSameName),
+                "non-matching instance must return false and leave registry unchanged");
+        assertSame(registered, toolkit.getTool("remove_if_same_wrong_ref"));
+    }
+
+    @Test
+    @DisplayName("removeToolIfSame returns false when deletion is disabled")
+    void removeToolIfSame_falseWhenDeletionDisabled() {
+        ToolkitConfig config = ToolkitConfig.builder().allowToolDeletion(false).build();
+        Toolkit tk = new Toolkit(config);
+        AgentTool tool = namedAgentTool("remove_if_same_no_delete");
+        tk.registerAgentTool(tool);
+        assertFalse(
+                tk.removeToolIfSame("remove_if_same_no_delete", tool),
+                "allowToolDeletion=false must return false");
+        assertSame(tool, tk.getTool("remove_if_same_no_delete"));
+    }
+
     @Test
     @DisplayName("Should ignore removeToolGroups when deletion is disabled")
     void testRemoveToolGroupsWhenDeletionDisabled() {