Skip to content

fix(session): Automatically detect and remove dangling tool calls whe… #873

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
liuxh613 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(session): Automatically detect and remove dangling tool calls whe… #873

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ee96e85
fix(session): Automatically detect and remove dangling tool calls whe…
liuxh613 Mar 5, 2026
dd21b04
ci: add ut
liuxh613 Mar 5, 2026
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
27 changes: 27 additions & 0 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 @@ -226,6 +226,8 @@ public void loadFrom(Session session, SessionKey sessionKey) {
// Load memory if managed
if (statePersistence.memoryManaged()) {
memory.loadFrom(session, sessionKey);
// Clean up any pending tool calls from previous session
cleanupPendingToolCalls();
}

// Load toolkit activeGroups if managed
Expand All @@ -240,6 +242,31 @@ public void loadFrom(Session session, SessionKey sessionKey) {
}
}

/**
* Clean up pending tool calls after session restoration.
*
* <p>When a session is restored, there may be incomplete tool calls from the previous
* session. This method removes the last assistant message if it contains tool calls without
* corresponding results, preventing IllegalStateException on the next user input.
*/
private void cleanupPendingToolCalls() {
Set<String> pendingIds = getPendingToolUseIds();
if (!pendingIds.isEmpty()) {
// Find and remove the last assistant message with pending tool calls
List<Msg> messages = memory.getMessages();
for (int i = messages.size() - 1; i >= 0; i--) {
Msg msg = messages.get(i);
if (msg.getRole() == MsgRole.ASSISTANT) {
memory.deleteMessage(i);
log.warn(
"Removed incomplete tool calls from restored session. Pending IDs: {}",
pendingIds);
break;
}
}
}
}

// ==================== Protected API ====================

@Override
Expand Down
221 changes: 221 additions & 0 deletions agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,9 @@
import io.agentscope.core.message.ToolUseBlock;
import io.agentscope.core.model.ChatResponse;
import io.agentscope.core.model.ChatUsage;
import io.agentscope.core.session.InMemorySession;
import io.agentscope.core.state.SimpleSessionKey;
import io.agentscope.core.state.StatePersistence;
import io.agentscope.core.tool.Toolkit;
import io.agentscope.core.util.JsonUtils;
import java.time.Duration;
Expand Down Expand Up @@ -932,6 +935,224 @@ reactor.core.publisher.Mono<T> onEvent(T event) {
"Second tool should be calculator");
}

@Test
@DisplayName("Test cleanup of pending tool calls after session restoration")
void testCleanupPendingToolCallsAfterSessionRestore() {
// Create a memory with incomplete tool calls (simulating restored session)
InMemoryMemory memory = new InMemoryMemory();

// Add user message
memory.addMessage(TestUtils.createUserMessage("User", "What's the weather?"));

// Add assistant message with tool call (but no result)
Msg assistantMsg =
Msg.builder()
.name(TestConstants.TEST_REACT_AGENT_NAME)
.role(MsgRole.ASSISTANT)
.content(
List.of(
ToolUseBlock.builder()
.name("weather")
.id("call_dangling_123")
.input(Map.of("city", "Beijing"))
.content("{\"city\":\"Beijing\"}")
.build()))
.build();
memory.addMessage(assistantMsg);

// Create mock model that will respond after cleanup
MockModel model = new MockModel("The weather is sunny.");

// Create agent with session persistence enabled
ReActAgent agent =
ReActAgent.builder()
.name(TestConstants.TEST_REACT_AGENT_NAME)
.sysPrompt(TestConstants.DEFAULT_SYS_PROMPT)
.model(model)
.memory(memory)
.statePersistence(StatePersistence.builder().memoryManaged(true).build())
.build();

// Simulate session restoration by calling loadFrom
InMemorySession session = new InMemorySession();
SimpleSessionKey sessionKey = SimpleSessionKey.of("test-session");

// Save memory to session first
memory.saveTo(session, sessionKey);

// Load from session (this should trigger cleanupPendingToolCalls)
agent.loadFrom(session, sessionKey);

// Verify that the dangling assistant message was removed
List<Msg> messages = agent.getMemory().getMessages();
assertEquals(1, messages.size(), "Should only have user message after cleanup");
assertEquals(
MsgRole.USER,
messages.get(0).getRole(),
"Remaining message should be user message");

// Verify agent can now accept new user input without error
Msg newUserMsg = TestUtils.createUserMessage("User", "Tell me more");
Msg response =
agent.call(newUserMsg)
.block(Duration.ofMillis(TestConstants.DEFAULT_TEST_TIMEOUT_MS));

assertNotNull(response, "Agent should respond successfully after cleanup");
assertEquals(
"The weather is sunny.", response.getTextContent(), "Response should match mock");
}

@Test
@DisplayName("Test no cleanup when no pending tool calls exist")
void testNoCleanupWhenNoPendingToolCalls() {
// Create a memory with complete tool call-result pair
InMemoryMemory memory = new InMemoryMemory();

// Add user message
memory.addMessage(TestUtils.createUserMessage("User", "What's the weather?"));

// Add assistant message with tool call
Msg assistantMsg =
Msg.builder()
.name(TestConstants.TEST_REACT_AGENT_NAME)
.role(MsgRole.ASSISTANT)
.content(
List.of(
ToolUseBlock.builder()
.name("weather")
.id("call_complete_456")
.input(Map.of("city", "Shanghai"))
.content("{\"city\":\"Shanghai\"}")
.build()))
.build();
memory.addMessage(assistantMsg);

// Add tool result
Msg toolResultMsg =
Msg.builder()
.name("User")
.role(MsgRole.USER)
.content(
List.of(
ToolResultBlock.builder()
.id("call_complete_456")
.output(
List.of(
TextBlock.builder()
.text("Sunny, 25°C")
.build()))
.build()))
.build();
memory.addMessage(toolResultMsg);

// Add final assistant response
Msg finalResponse =
Msg.builder()
.name(TestConstants.TEST_REACT_AGENT_NAME)
.role(MsgRole.ASSISTANT)
.content(
List.of(
TextBlock.builder()
.text("The weather in Shanghai is sunny, 25°C.")
.build()))
.build();
memory.addMessage(finalResponse);

// Create agent with session persistence
ReActAgent agent =
ReActAgent.builder()
.name(TestConstants.TEST_REACT_AGENT_NAME)
.sysPrompt(TestConstants.DEFAULT_SYS_PROMPT)
.model(new MockModel(TestConstants.MOCK_MODEL_SIMPLE_RESPONSE))
.memory(memory)
.statePersistence(StatePersistence.builder().memoryManaged(true).build())
.build();

// Simulate session restoration
InMemorySession session = new InMemorySession();
SimpleSessionKey sessionKey = SimpleSessionKey.of("test-session-2");

memory.saveTo(session, sessionKey);
agent.loadFrom(session, sessionKey);

// Verify that no messages were removed (all 4 messages should remain)
List<Msg> messages = agent.getMemory().getMessages();
assertEquals(4, messages.size(), "All messages should remain when no pending tool calls");
assertEquals(MsgRole.USER, messages.get(0).getRole(), "First message should be user");
assertEquals(
MsgRole.ASSISTANT, messages.get(1).getRole(), "Second message should be assistant");
assertEquals(
MsgRole.USER,
messages.get(2).getRole(),
"Third message should be user (tool result)");
assertEquals(
MsgRole.ASSISTANT, messages.get(3).getRole(), "Fourth message should be assistant");
}

@Test
@DisplayName("Test cleanup with multiple assistant messages")
void testCleanupOnlyLastAssistantMessage() {
// Create a memory with multiple assistant messages, last one has pending tool call
InMemoryMemory memory = new InMemoryMemory();

// First conversation turn (complete)
memory.addMessage(TestUtils.createUserMessage("User", "Hello"));
memory.addMessage(
Msg.builder()
.name(TestConstants.TEST_REACT_AGENT_NAME)
.role(MsgRole.ASSISTANT)
.content(List.of(TextBlock.builder().text("Hi there!").build()))
.build());

// Second conversation turn (incomplete - has pending tool call)
memory.addMessage(TestUtils.createUserMessage("User", "What's the weather?"));
Msg pendingAssistantMsg =
Msg.builder()
.name(TestConstants.TEST_REACT_AGENT_NAME)
.role(MsgRole.ASSISTANT)
.content(
List.of(
ToolUseBlock.builder()
.name("weather")
.id("call_pending_789")
.input(Map.of("city", "Guangzhou"))
.content("{\"city\":\"Guangzhou\"}")
.build()))
.build();
memory.addMessage(pendingAssistantMsg);

// Create agent
ReActAgent agent =
ReActAgent.builder()
.name(TestConstants.TEST_REACT_AGENT_NAME)
.sysPrompt(TestConstants.DEFAULT_SYS_PROMPT)
.model(new MockModel(TestConstants.MOCK_MODEL_SIMPLE_RESPONSE))
.memory(memory)
.statePersistence(StatePersistence.builder().memoryManaged(true).build())
.build();

// Simulate session restoration
InMemorySession session = new InMemorySession();
SimpleSessionKey sessionKey = SimpleSessionKey.of("test-session-3");

memory.saveTo(session, sessionKey);
agent.loadFrom(session, sessionKey);

// Verify that only the last assistant message was removed
List<Msg> messages = agent.getMemory().getMessages();
assertEquals(3, messages.size(), "Should have 3 messages after cleanup");

// Verify the remaining messages
assertEquals(MsgRole.USER, messages.get(0).getRole(), "First message should be user");
assertEquals(
MsgRole.ASSISTANT, messages.get(1).getRole(), "Second message should be assistant");
assertEquals(
"Hi there!",
messages.get(1).getTextContent(),
"Second message content should be preserved");
assertEquals(MsgRole.USER, messages.get(2).getRole(), "Third message should be user");
}

// Helper method to create tool call response
private static ChatResponse createToolCallResponseHelper(
String toolName, String toolCallId, Map<String, Object> arguments) {
Expand Down
Loading