From a94a756677864c4d00dbccdc28cf40714c550619 Mon Sep 17 00:00:00 2001 From: matevip Date: Tue, 23 Jun 2026 18:22:56 +0800 Subject: [PATCH] feat(agent): register send continuations in the sub-agent registry --- .../mate/tool/builtin/SessionSendTool.java | 8 ++++- .../tool/builtin/SessionSendToolTest.java | 29 ++++++++++++++++--- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/mateclaw-server/src/main/java/vip/mate/tool/builtin/SessionSendTool.java b/mateclaw-server/src/main/java/vip/mate/tool/builtin/SessionSendTool.java index af1b2559..363aef20 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/builtin/SessionSendTool.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/builtin/SessionSendTool.java @@ -10,6 +10,7 @@ import org.springframework.lang.Nullable; import org.springframework.stereotype.Component; import vip.mate.agent.AgentService; import vip.mate.agent.context.ChatOrigin; +import vip.mate.agent.delegation.SubagentRegistry; import vip.mate.workspace.conversation.model.ConversationEntity; import vip.mate.workspace.conversation.repository.ConversationMapper; @@ -40,6 +41,7 @@ public class SessionSendTool { private final AgentService agentService; private final ConversationMapper conversationMapper; + private final SubagentRegistry subagentRegistry; @Tool(description = """ Send a follow-up message to a sub-agent you previously delegated to, continuing its @@ -90,9 +92,12 @@ public class SessionSendTool { Long agentId = child.getAgentId(); ChatOrigin origin = ChatOrigin.from(ctx).withAgent(agentId).withConversationId(sessionId); - // Re-enter the delegation context one level below the caller so the + // Register the continuation so an in-flight follow-up is visible to + // SessionListTool and interruptible via the subagent control API, then + // re-enter the delegation context one level below the caller so the // continued child stays gated (cannot delegate / send onward) and the // depth cap keeps holding for anything it tries to spawn. + String subagentId = subagentRegistry.register(callerConversationId, sessionId, agentId, message, null); DelegationContext.enter(callerConversationId, DelegateAgentTool.DEFAULT_CHILD_DENIED_TOOLS, resolveRootConversationId(callerConversationId), null, callerDepth + 1); try { @@ -104,6 +109,7 @@ public class SessionSendTool { return "[Error] Sub-agent follow-up failed: " + e.getMessage(); } finally { DelegationContext.exit(); + subagentRegistry.unregister(subagentId); } } diff --git a/mateclaw-server/src/test/java/vip/mate/tool/builtin/SessionSendToolTest.java b/mateclaw-server/src/test/java/vip/mate/tool/builtin/SessionSendToolTest.java index 25870853..00b6bf70 100644 --- a/mateclaw-server/src/test/java/vip/mate/tool/builtin/SessionSendToolTest.java +++ b/mateclaw-server/src/test/java/vip/mate/tool/builtin/SessionSendToolTest.java @@ -7,14 +7,15 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.InjectMocks; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import vip.mate.agent.AgentService; import vip.mate.agent.context.ChatOrigin; +import vip.mate.agent.delegation.SubagentRegistry; import vip.mate.workspace.conversation.model.ConversationEntity; import vip.mate.workspace.conversation.repository.ConversationMapper; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; @@ -34,7 +35,8 @@ class SessionSendToolTest { @Mock AgentService agentService; @Mock ConversationMapper conversationMapper; - @InjectMocks SessionSendTool tool; + private SubagentRegistry registry; + private SessionSendTool tool; @BeforeAll static void initMyBatisPlusCache() { @@ -45,6 +47,8 @@ class SessionSendToolTest { @BeforeEach void setUp() { + registry = new SubagentRegistry(); + tool = new SessionSendTool(agentService, conversationMapper, registry); ToolExecutionContext.clear(); while (DelegationContext.currentDepth() > 0) { DelegationContext.exit(); @@ -88,7 +92,6 @@ class SessionSendToolTest { @Test void rejectsSessionOwnedByAnotherConversation() { ToolExecutionContext.set("conv-root", "tester"); - // Child's parent is a different conversation than the caller. when(conversationMapper.selectOne(any())).thenReturn(child("child-1", "other-conv", 7L)); String out = tool.sendToSubagent("child-1", "do more", null); assertTrue(out.contains("does not belong to this conversation"), out); @@ -98,7 +101,6 @@ class SessionSendToolTest { @Test void rejectsWhenDepthLimitReached() { ToolExecutionContext.set("conv-root", "tester"); - // Simulate being already at the max delegation depth. DelegationContext.enter("conv", java.util.Set.of(), "root", "sa", DelegateAgentTool.MAX_DELEGATION_DEPTH); try { String out = tool.sendToSubagent("child-1", "do more", null); @@ -123,4 +125,23 @@ class SessionSendToolTest { assertTrue(out.contains("refined result"), out); verify(agentService).chat(eq(7L), eq("refine it"), eq("child-1"), any(ChatOrigin.class)); } + + @Test + void registersDuringContinuationAndUnregistersAfter() { + ToolExecutionContext.set("conv-root", "tester"); + when(conversationMapper.selectOne(any())).thenReturn(child("child-1", "conv-root", 7L)); + // While the child runs, the continuation must be visible in the registry + // (so SessionListTool / the control API can see and interrupt it). + when(agentService.chat(eq(7L), any(), eq("child-1"), any(ChatOrigin.class))).thenAnswer(inv -> { + assertFalse(registry.snapshot("conv-root").isEmpty(), + "continuation should be registered while running"); + return "done"; + }); + + tool.sendToSubagent("child-1", "keep going", null); + + // ...and cleaned up afterwards so it never leaks. + assertTrue(registry.snapshot("conv-root").isEmpty(), + "continuation should be unregistered after completion"); + } }