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 2fd0eef59a..9a5a0088eb 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java +++ b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java @@ -486,24 +486,49 @@ private Mono saveStateToSession(CallExecution scope) { if (stateStore == null) { return Mono.empty(); } + return Mono.fromRunnable(() -> persistState(scope)) + .subscribeOn(Schedulers.boundedElastic()); + } + + private void persistState(CallExecution scope) { syncToolkitToState(scope.state); SlotRef ref = SlotRef.parse(scope.slotKey); - AgentState toSave = scope.state; - return Mono.fromRunnable( - () -> { - long newVersion = - persistAgentStateCas( - ref.userId, - ref.sessionId, - scope.slotKey, - toSave, - scope.loadedVersion, - scope.loadedContextSize); - if (newVersion != AgentStateStore.UNVERSIONED) { - scope.loadedVersion = newVersion; - } - }) - .subscribeOn(Schedulers.boundedElastic()); + long newVersion = + persistAgentStateCas( + ref.userId, + ref.sessionId, + scope.slotKey, + scope.state, + scope.loadedVersion, + scope.loadedContextSize); + if (newVersion != AgentStateStore.UNVERSIONED) { + scope.loadedVersion = newVersion; + } + } + + private void repairStateBeforeAbnormalCheckpoint(CallExecution scope) { + try { + scope.dropUncommittedToolCallsFromCurrentCall(); + } catch (RuntimeException repairError) { + log.warn("Failed to repair agent state before checkpoint", repairError); + } + } + + private void checkpointStateAfterCancellation(CallExecution scope) { + repairStateBeforeAbnormalCheckpoint(scope); + try { + if (stateStore != null) { + persistState(scope); + } + } catch (RuntimeException saveError) { + log.warn("Failed to save agent state after abnormal termination", saveError); + } + } + + private Mono checkpointOnAbnormalTermination(CallExecution scope, Mono execution) { + return execution + .doOnCancel(() -> checkpointStateAfterCancellation(scope)) + .onErrorResume(error -> saveStateAfterCallFailure(scope, error)); } /** @@ -522,6 +547,7 @@ private Mono saveStateAfterCallFailure(CallExecution scope, Throwable cal if (ExceptionUtils.containsInterruptedException(callFailure)) { return Mono.error(callFailure); } + repairStateBeforeAbnormalCheckpoint(scope); return saveStateToSession(scope) .onErrorResume( saveFailure -> { @@ -1223,8 +1249,7 @@ protected Mono doCall(List msgs) { AgentEventEmitter.fromForwardingContext(cv) .ifPresent(ae -> scope.externalEventEmitter = ae); } - return scope.doCallInner(msgs) - .onErrorResume(error -> saveStateAfterCallFailure(scope, error)) + return checkpointOnAbnormalTermination(scope, scope.doCallInner(msgs)) .flatMap(result -> saveStateToSession(scope).thenReturn(result)); }); } @@ -1265,20 +1290,25 @@ private Mono doStructuredCall(List msgs, Class targetClass, JsonNod hasTools ? model.supportsNativeStructuredOutputWithTools() : model.supportsNativeStructuredOutput(); + Mono execution; 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); - }); + execution = + 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); + }); + } else { + execution = doFallbackStructuredCall(msgs, jsonSchema); } - return doFallbackStructuredCall(msgs, jsonSchema); + return Mono.deferContextual( + cv -> checkpointOnAbnormalTermination(scopeFrom(cv), execution)); } /** @@ -1359,7 +1389,6 @@ private Mono doFallbackStructuredCall(List msgs, Map j scope.soTool = createStructuredOutputTool(jsonSchema); return scope.doCallInner(msgs) - .onErrorResume(error -> saveStateAfterCallFailure(scope, error)) .flatMap( result -> { Msg out = result; @@ -2124,6 +2153,40 @@ private void synthesizeErrorResultsForPendingToolCalls() { } } + /** Remove current-call PENDING tool calls that never produced results. */ + private void dropUncommittedToolCallsFromCurrentCall() { + Set pendingIds = getPendingToolUseIds(); + if (pendingIds.isEmpty()) { + return; + } + List context = state.contextMutable(); + for (int i = context.size() - 1; i >= loadedContextSize; i--) { + Msg message = context.get(i); + if (message.getRole() != MsgRole.ASSISTANT + || !message.hasContentBlocks(ToolUseBlock.class)) { + continue; + } + List retained = + message.getContent().stream() + .filter( + block -> + !(block instanceof ToolUseBlock toolUse) + || toolUse.getState() + != ToolCallState.PENDING + || !pendingIds.contains(toolUse.getId())) + .toList(); + if (retained.size() == message.getContent().size()) { + return; + } + if (retained.isEmpty()) { + context.remove(i); + } else { + context.set(i, message.withContent(retained)); + } + return; + } + } + private void publishEvent(AgentEvent event) { FluxSink sink = eventSink; if (sink != null) { 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 0bb4fc0554..7e14c124bd 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 @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotSame; 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 io.agentscope.core.ReActAgent; @@ -30,6 +31,9 @@ import io.agentscope.core.message.Msg; import io.agentscope.core.message.MsgRole; import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ToolUseBlock; +import io.agentscope.core.middleware.ActingInput; +import io.agentscope.core.middleware.MiddlewareBase; import io.agentscope.core.model.ChatModelBase; import io.agentscope.core.model.ChatResponse; import io.agentscope.core.model.GenerateOptions; @@ -51,11 +55,13 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; @@ -64,6 +70,10 @@ @DisplayName("ReActAgent per-session state API") class ReActAgentPerSessionStateTest { + static class StructuredResponse { + public String value; + } + private static final class NoopModel extends ChatModelBase { @Override public String getModelName() { @@ -435,6 +445,141 @@ void userInterruptPersistsRecoveryState() throws Exception { assertEquals(GenerateReason.INTERRUPTED, restoredRecovery.getGenerateReason()); } + @Test + @DisplayName("model errors checkpoint committed context without partial output") + void modelErrorCheckpointsCommittedContext() { + InMemoryAgentStateStore store = new InMemoryAgentStateStore(); + RuntimeContext ctx = RuntimeContext.builder().userId("u1").sessionId("model-error").build(); + ReActAgent first = + ReActAgent.builder() + .name("asst") + .sysPrompt("hi") + .model(failingReasoningModel()) + .stateStore(store) + .build(); + IllegalStateException error = + assertThrows( + IllegalStateException.class, + () -> + first.streamEvents(List.of(userMsg("hello")), ctx) + .blockLast(Duration.ofSeconds(5))); + + assertEquals("model stream failed", error.getMessage()); + AgentState restored = agent(store).getAgentState("u1", "model-error"); + assertTrue(allText(restored).contains("hello")); + assertFalse(allText(restored).contains("visible before error")); + assertFalse(hasToolUse(restored, "call-malformed")); + } + + @Test + @DisplayName("acting errors remove newly persisted tool calls without results") + void actingErrorDropsUncommittedToolCallBeforeCheckpoint() { + InMemoryAgentStateStore store = new InMemoryAgentStateStore(); + RuntimeContext ctx = + RuntimeContext.builder().userId("u1").sessionId("acting-error").build(); + ReActAgent first = + ReActAgent.builder() + .name("asst") + .sysPrompt("hi") + .model(completedToolCallModel()) + .stateStore(store) + .middleware(new FailingActingMiddleware()) + .build(); + + IllegalStateException error = + assertThrows( + IllegalStateException.class, + () -> + first.streamEvents(List.of(userMsg("hello")), ctx) + .blockLast(Duration.ofSeconds(5))); + + assertEquals("acting failed", error.getMessage()); + AgentState restored = agent(store).getAgentState("u1", "acting-error"); + assertTrue(allText(restored).contains("completed response")); + assertFalse(hasToolUse(restored, "call-complete")); + } + + @Test + @DisplayName("stream cancellation checkpoints committed context") + void cancellationCheckpointsCommittedContext() throws Exception { + InMemoryAgentStateStore store = new InMemoryAgentStateStore(); + CountDownLatch chunksEmitted = new CountDownLatch(1); + ChatModelBase model = + model( + Flux.just( + response( + TextBlock.builder() + .text("visible before cancel") + .build()), + response( + ToolUseBlock.builder() + .id("call-cancelled") + .name("echo") + .content("{\"value\":") + .build())) + .doOnComplete(chunksEmitted::countDown) + .concatWith(Flux.never())); + RuntimeContext ctx = RuntimeContext.builder().userId("u1").sessionId("cancelled").build(); + ReActAgent first = + ReActAgent.builder() + .name("asst") + .sysPrompt("hi") + .model(model) + .stateStore(store) + .build(); + + Disposable subscription = + first.streamEvents(List.of(userMsg("hello")), ctx).subscribe(event -> {}); + assertTrue( + chunksEmitted.await(5, TimeUnit.SECONDS), + "visible chunks should be consumed before cancellation"); + + subscription.dispose(); + + AgentState restored = agent(store).getAgentState("u1", "cancelled"); + assertTrue(allText(restored).contains("hello")); + assertFalse(allText(restored).contains("visible before cancel")); + assertFalse(hasToolUse(restored, "call-cancelled")); + } + + @Test + @DisplayName("native structured-output cancellation checkpoints committed context") + void nativeStructuredOutputCancellationCheckpointsCommittedContext() throws Exception { + InMemoryAgentStateStore store = new InMemoryAgentStateStore(); + CountDownLatch chunkEmitted = new CountDownLatch(1); + ChatModelBase model = + model( + Flux.just( + response( + TextBlock.builder() + .text("visible before cancel") + .build())) + .doOnComplete(chunkEmitted::countDown) + .concatWith(Flux.never()), + true); + RuntimeContext ctx = + RuntimeContext.builder().userId("u1").sessionId("structured-cancel").build(); + ReActAgent first = + ReActAgent.builder() + .name("asst") + .sysPrompt("hi") + .model(model) + .stateStore(store) + .build(); + + Disposable subscription = + first.call(List.of(userMsg("hello")), StructuredResponse.class, ctx).subscribe(); + assertTrue( + chunkEmitted.await(5, TimeUnit.SECONDS), + "visible chunk should be consumed before cancellation"); + + subscription.dispose(); + + AgentState restored = agent(store).getAgentState("u1", "structured-cancel"); + assertTrue(allText(restored).contains("hello")); + assertFalse(allText(restored).contains("visible before cancel")); + } + private static final class DelayedFirstChunkModel extends ChatModelBase { private final CountDownLatch subscribed; @@ -466,6 +611,72 @@ protected Flux doStream( } } + private static ChatModelBase completedToolCallModel() { + return model( + Flux.just( + response( + TextBlock.builder().text("completed response").build(), + ToolUseBlock.builder() + .id("call-complete") + .name("echo") + .input(java.util.Map.of("value", "done")) + .build()))); + } + + private static ChatModelBase failingReasoningModel() { + return model( + Flux.concat( + Flux.just( + response(TextBlock.builder().text("visible before error").build()), + response( + ToolUseBlock.builder() + .id("call-malformed") + .name("echo") + .content("{\"value\":") + .build())), + Flux.error(new IllegalStateException("model stream failed")))); + } + + private static ChatModelBase model(Flux responses) { + return model(responses, false); + } + + private static ChatModelBase model( + Flux responses, boolean supportsNativeStructuredOutput) { + return new ChatModelBase() { + @Override + public String getModelName() { + return "test"; + } + + @Override + protected Flux doStream( + List messages, List tools, GenerateOptions options) { + return responses; + } + + @Override + public boolean supportsNativeStructuredOutput() { + return supportsNativeStructuredOutput; + } + }; + } + + private static final class FailingActingMiddleware implements MiddlewareBase { + @Override + public Flux onActing( + Agent agent, + RuntimeContext ctx, + ActingInput input, + Function> next) { + return Flux.error(new IllegalStateException("acting failed")); + } + } + + private static ChatResponse response(ContentBlock... blocks) { + return ChatResponse.builder().content(List.of(blocks)).build(); + } + private static Msg userMsg(String text) { return Msg.builder() .name("user") @@ -474,6 +685,12 @@ private static Msg userMsg(String text) { .build(); } + private static boolean hasToolUse(AgentState state, String id) { + return state.getContext().stream() + .flatMap(msg -> msg.getContentBlocks(ToolUseBlock.class).stream()) + .anyMatch(block -> id.equals(block.getId())); + } + private static List allText(AgentState state) { List out = new ArrayList<>(); for (Msg m : state.getContext()) {