From 67404d5410ed451a8b3734b668d86272532554e3 Mon Sep 17 00:00:00 2001 From: Sparkle6979 <111422548+Sparkle6979@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:28:16 +0800 Subject: [PATCH] Offload agent execution setup from caller thread --- .../io/agentscope/core/agent/AgentBase.java | 53 ++++++--- .../agentscope/core/agent/AgentBaseTest.java | 105 ++++++++++++++++++ 2 files changed, 140 insertions(+), 18 deletions(-) diff --git a/agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java b/agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java index 72e5d82045..8d10497a6a 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java +++ b/agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java @@ -243,6 +243,8 @@ protected Mono callInternal( public static final String SHUTDOWN_REQUEST_ID_KEY = "io.agentscope.core.agent.AgentBase.shutdownRequestId"; + private record PreparedExecution(Object scope) {} + /** * Shared {@code call()} lifecycle: acquire execution, then (inside {@code deferContextual} so * the caller-supplied {@link RuntimeContext} is read per-subscription) run {@link @@ -302,24 +304,39 @@ private Mono runLifecycleBody( RuntimeContext rc, Function, Mono> doCallFn, String requestId) { - Object scope = beforeAgentExecution(msgs, rc); - // Bind this call's resolved per-session state to the tracked shutdown request so graceful - // shutdown interrupts / saves the exact (userId, sessionId) session rather than the agent's - // no-arg "most-recently-active" accessors. - GracefulShutdownManager.getInstance().bindRequestState(requestId, stateForCall(scope)); - Mono body = - TracerRegistry.get() - .callAgent( - this, - msgs, - () -> - notifyPreCall(msgs, scope) - .flatMap(doCallFn) - .flatMap(this::notifyPostCall) - .onErrorResume( - createErrorHandler( - msgs.toArray(new Msg[0])))); - return scope == null ? body : body.contextWrite(c -> c.put(CALL_SCOPE_KEY, scope)); + return Mono.fromCallable( + () -> { + Object scope = beforeAgentExecution(msgs, rc); + // Bind this call's resolved per-session state to the tracked shutdown + // request so graceful shutdown interrupts / saves the exact (userId, + // sessionId) session rather than the agent's no-arg + // "most-recently-active" accessors. + GracefulShutdownManager.getInstance() + .bindRequestState(requestId, stateForCall(scope)); + return new PreparedExecution(scope); + }) + .subscribeOn(Schedulers.boundedElastic()) + .flatMap( + prepared -> { + Object scope = prepared.scope(); + Mono body = + TracerRegistry.get() + .callAgent( + this, + msgs, + () -> + notifyPreCall(msgs, scope) + .flatMap(doCallFn) + .flatMap(this::notifyPostCall) + .onErrorResume( + createErrorHandler( + msgs.toArray( + new Msg + [0])))); + return scope == null + ? body + : body.contextWrite(c -> c.put(CALL_SCOPE_KEY, scope)); + }); } /** diff --git a/agentscope-core/src/test/java/io/agentscope/core/agent/AgentBaseTest.java b/agentscope-core/src/test/java/io/agentscope/core/agent/AgentBaseTest.java index 9e2919643d..7e503428be 100644 --- a/agentscope-core/src/test/java/io/agentscope/core/agent/AgentBaseTest.java +++ b/agentscope-core/src/test/java/io/agentscope/core/agent/AgentBaseTest.java @@ -30,6 +30,8 @@ import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; @@ -110,6 +112,63 @@ protected Mono handleInterrupt(InterruptContext context, Msg... originalArg } } + static class LifecycleTestAgent extends TestAgent { + private final AtomicReference beforeThreadName = new AtomicReference<>(); + private final AtomicInteger activeBeforeCalls = new AtomicInteger(); + private final AtomicInteger maxActiveBeforeCalls = new AtomicInteger(); + private final AtomicBoolean firstBeforeEntered = new AtomicBoolean(); + private Duration beforeDelay = Duration.ZERO; + private Object serializationKey; + + LifecycleTestAgent(String name) { + super(name); + } + + void setBeforeDelay(Duration beforeDelay) { + this.beforeDelay = beforeDelay; + } + + void setSerializationKey(Object serializationKey) { + this.serializationKey = serializationKey; + } + + String getBeforeThreadName() { + return beforeThreadName.get(); + } + + int getMaxActiveBeforeCalls() { + return maxActiveBeforeCalls.get(); + } + + boolean hasEnteredBefore() { + return firstBeforeEntered.get(); + } + + @Override + protected Object beforeAgentExecution(List msgs, RuntimeContext rc) { + beforeThreadName.set(Thread.currentThread().getName()); + firstBeforeEntered.set(true); + int active = activeBeforeCalls.incrementAndGet(); + maxActiveBeforeCalls.accumulateAndGet(active, Math::max); + try { + if (!beforeDelay.isZero()) { + Thread.sleep(beforeDelay.toMillis()); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } finally { + activeBeforeCalls.decrementAndGet(); + } + return null; + } + + @Override + protected Object callSerializationKey(RuntimeContext rc) { + return serializationKey; + } + } + @BeforeEach void setUp() { agent = new TestAgent(TestConstants.TEST_AGENT_NAME); @@ -200,6 +259,52 @@ void testConcurrencyConflict() { assertNotNull(response2, "Response should not be null"); } + @Test + @DisplayName("Should run beforeAgentExecution on boundedElastic and keep null scope valid") + void testBeforeAgentExecutionRunsOnBoundedElastic() { + LifecycleTestAgent lifecycleAgent = new LifecycleTestAgent("LifecycleAgent"); + + Msg response = + lifecycleAgent + .call(TestUtils.createUserMessage("User", "thread check")) + .subscribeOn(Schedulers.single()) + .block(Duration.ofMillis(TestConstants.DEFAULT_TEST_TIMEOUT_MS)); + + assertNotNull( + response, "Response should not be null when beforeAgentExecution returns null"); + assertTrue(lifecycleAgent.hasEnteredBefore(), "beforeAgentExecution should run"); + assertTrue( + lifecycleAgent.getBeforeThreadName().contains("boundedElastic"), + "beforeAgentExecution should run on boundedElastic, but ran on " + + lifecycleAgent.getBeforeThreadName()); + } + + @Test + @DisplayName("Should keep serialized beforeAgentExecution calls for the same key") + void testBeforeAgentExecutionStillSerializedByKey() { + LifecycleTestAgent lifecycleAgent = new LifecycleTestAgent("SerializedLifecycleAgent"); + lifecycleAgent.setSerializationKey("same-session"); + lifecycleAgent.setBeforeDelay(Duration.ofMillis(150)); + + Msg msg1 = TestUtils.createUserMessage("User", "call 1"); + Msg msg2 = TestUtils.createUserMessage("User", "call 2"); + + List results = + Mono.zip( + lifecycleAgent.call(msg1).subscribeOn(Schedulers.parallel()), + lifecycleAgent.call(msg2).subscribeOn(Schedulers.parallel())) + .map(tuple -> List.of(tuple.getT1(), tuple.getT2())) + .block(Duration.ofSeconds(5)); + + assertNotNull(results, "Both calls should complete"); + assertEquals(2, results.size(), "Both calls should return responses"); + assertEquals( + 1, + lifecycleAgent.getMaxActiveBeforeCalls(), + "beforeAgentExecution should not overlap for calls with the same serialization" + + " key"); + } + @Test @DisplayName("Should handle observe without generating reply") void testObserve() {