From 81c6488a3c107d8bdeb0da6fb8e3589b9b74f5b4 Mon Sep 17 00:00:00 2001 From: matevip Date: Sat, 16 May 2026 14:51:20 +0800 Subject: [PATCH] feat(tool): async sub-agent delegation with attribution gate --- .../mate/tool/builtin/DelegateAgentTool.java | 300 +++++++++++++++ .../DelegateAgentToolDenyListTest.java | 3 +- ...elegateAsyncTaskOutputAttributionTest.java | 203 ++++++++++ .../tool/builtin/DelegateAsyncToolTest.java | 347 ++++++++++++++++++ 4 files changed, 852 insertions(+), 1 deletion(-) create mode 100644 mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAsyncTaskOutputAttributionTest.java create mode 100644 mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAsyncToolTest.java diff --git a/mateclaw-server/src/main/java/vip/mate/tool/builtin/DelegateAgentTool.java b/mateclaw-server/src/main/java/vip/mate/tool/builtin/DelegateAgentTool.java index 4fcdd294..0159c6de 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/builtin/DelegateAgentTool.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/builtin/DelegateAgentTool.java @@ -2,6 +2,7 @@ package vip.mate.tool.builtin; import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper; import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -18,8 +19,11 @@ import vip.mate.agent.model.AgentEntity; import vip.mate.agent.repository.AgentMapper; import vip.mate.audit.service.AuditEventService; import vip.mate.channel.web.ChatStreamTracker; +import vip.mate.task.AsyncTaskService; +import vip.mate.task.model.AsyncTaskEntity; import vip.mate.workspace.conversation.ConversationService; +import java.time.Duration; import java.util.*; import java.util.concurrent.*; import java.util.stream.Collectors; @@ -128,6 +132,27 @@ public class DelegateAgentTool { private final ObjectMapper objectMapper; private final SubagentRegistry subagentRegistry; private final AuditEventService auditEventService; + private final AsyncTaskService asyncTaskService; + + /** Max characters of the task description persisted in {@code request_json}. + * Anything longer is truncated — full task is still inside the running + * child's conversation context. */ + private static final int ASYNC_TASK_REQUEST_MAX_CHARS = 8000; + + /** Max label length carried inside {@code request_json} and surfaced on + * spawn-event payloads. Picked to fit a short UI badge without wrapping. */ + private static final int ASYNC_LABEL_MAX_CHARS = 32; + + /** Default {@code block=true} wait when caller omits {@code timeoutSeconds}. */ + private static final int TASK_OUTPUT_DEFAULT_TIMEOUT_S = 30; + + /** Upper bound on {@code block=true} wait. Picked to be longer than the + * typical ReAct turn latency yet short enough that the parent agent + * doesn't burn its own LLM budget blocked on a stalled child. */ + private static final int TASK_OUTPUT_MAX_TIMEOUT_S = 120; + + /** Polling interval inside {@code block=true} wait. */ + private static final long TASK_OUTPUT_POLL_INTERVAL_MS = 500L; /** * Operator-supplied deny-list extension. Configured via @@ -565,6 +590,281 @@ public class DelegateAgentTool { return truncate(sb.toString(), MAX_RESULT_LENGTH * 2); // 并行结果允许更长 } + // ==================== Async (detached) delegation ==================== + + @Tool(description = """ + Delegate a task to another agent asynchronously and return a task_id immediately. \ + Parent continues reasoning while child runs in background. \ + Use task_output(task_id) in a later turn to retrieve the result. \ + Best for long-running sub-tasks (research, file processing) where the parent has \ + other work to do in parallel. For quick tasks where you need the answer immediately, \ + use delegateToAgent instead.""") + public String delegateAsync( + @ToolParam(description = "Target Agent name (exact match)") String agentName, + @ToolParam(description = "Task description with complete context information") String task, + @ToolParam(description = "Optional short label (≤ 32 chars) for human tracking on the UI badge", + required = false) String label, + @Nullable ToolContext ctx) { + + if (agentName == null || agentName.isBlank()) { + return errorJson("agentName 不能为空"); + } + if (task == null || task.isBlank()) { + return errorJson("task 不能为空"); + } + String safeLabel = label == null ? "" : + (label.length() > ASYNC_LABEL_MAX_CHARS ? label.substring(0, ASYNC_LABEL_MAX_CHARS) : label); + + int depth = DelegationContext.currentDepth(); + if (depth >= MAX_DELEGATION_DEPTH) { + return errorJson("Delegation depth exceeded (max " + MAX_DELEGATION_DEPTH + ")"); + } + + AgentEntity target = findAgent(agentName); + if (target == null) { + return errorJson("Agent not found: " + agentName); + } + + String parentConversationId = resolveParentConversationId(); + if (parentConversationId == null || parentConversationId.isBlank()) { + return errorJson("delegateAsync requires a parent conversation context"); + } + if (subagentRegistry.isSpawnPaused(parentConversationId)) { + return errorJson("Spawning paused for this conversation; resume via /api/v1/subagents/spawn-pause"); + } + + // Capture origin / user on the calling thread — the Callable runs on + // AsyncTaskService.pollExecutor, where the ToolContext ThreadLocal is + // not visible. The child's identity (agentId) is swapped in below; + // channel / workspace / requester all propagate via the closure. + ChatOrigin parentOrigin = ChatOrigin.from(ctx); + String currentUser = parentOrigin != null && parentOrigin.requesterId() != null + && !parentOrigin.requesterId().isBlank() + ? parentOrigin.requesterId() + : "system"; + + String childConversationId = createChildConv(target, parentConversationId); + + String requestJson; + try { + Map payload = new LinkedHashMap<>(); + payload.put("parentConversationId", parentConversationId); + payload.put("childConversationId", childConversationId); + payload.put("childAgentId", target.getId()); + payload.put("task", truncate(task, ASYNC_TASK_REQUEST_MAX_CHARS)); + payload.put("label", safeLabel); + requestJson = objectMapper.writeValueAsString(payload); + } catch (Exception e) { + return errorJson("Failed to serialize task payload: " + e.getMessage()); + } + + // Live observability handle — task_output never reads from it; the + // persistent mate_async_task row is the source of truth for status, + // result, and attribution. + String subagentId = subagentRegistry.register(parentConversationId, childConversationId, + target.getId(), task, null); + + AsyncTaskEntity entity; + try { + entity = asyncTaskService.submitOneShot( + "agent_delegate", + parentConversationId, + null, + requestJson, + currentUser, + () -> { + try { + ChildResult childResult = runSingleChild(0, target, task, + parentConversationId, childConversationId, parentOrigin); + return childResult.toToolResponse(target.getName()); + } finally { + subagentRegistry.get(subagentId).ifPresent(rec -> { + if ("running".equals(rec.status().get())) { + rec.status().set("completed"); + } + }); + subagentRegistry.unregister(subagentId); + } + }); + } catch (IllegalStateException e) { + // Per-user concurrency cap hit inside AsyncTaskService#createTask. + // Roll back the registry entry so it doesn't dangle. + subagentRegistry.unregister(subagentId); + return errorJson(e.getMessage()); + } catch (Exception e) { + subagentRegistry.unregister(subagentId); + log.error("delegateAsync submit failed: target={}, err={}", target.getName(), e.getMessage()); + return errorJson("Failed to spawn async task: " + e.getMessage()); + } + + log.info("Async delegation spawned: taskId={}, target={}({}), childConv={}, parentConv={}", + entity.getTaskId(), target.getName(), target.getId(), + childConversationId, parentConversationId); + + if (streamTracker.isRunning(parentConversationId)) { + Map spawnEvent = new LinkedHashMap<>(); + spawnEvent.put("taskId", entity.getTaskId()); + spawnEvent.put("childConversationId", childConversationId); + spawnEvent.put("childAgentName", target.getName()); + spawnEvent.put("label", safeLabel); + spawnEvent.put("task", truncate(task, 200)); + streamTracker.broadcastObject(parentConversationId, "delegation_async_spawned", spawnEvent); + } + + Map result = new LinkedHashMap<>(); + result.put("task_id", entity.getTaskId()); + result.put("child_conversation_id", childConversationId); + result.put("agent_name", target.getName()); + result.put("status", "running"); + result.put("hint", "Call task_output(task_id) in a later turn to retrieve the result."); + if (!safeLabel.isEmpty()) { + result.put("label", safeLabel); + } + try { + return objectMapper.writeValueAsString(result); + } catch (Exception e) { + return errorJson("Failed to serialize response: " + e.getMessage()); + } + } + + @Tool(description = """ + Retrieve the result of a previously spawned async sub-agent task. \ + Returns the final reply when completed, or a status indicator if still running. \ + Set block=true to wait up to timeout seconds for completion.""") + public String taskOutput( + @ToolParam(description = "task_id returned by delegateAsync") String taskId, + @ToolParam(description = "Whether to block until done or timeout. Default false.", + required = false) Boolean block, + @ToolParam(description = "Max seconds to wait when block=true. Default 30, max 120.", + required = false) Integer timeoutSeconds, + @Nullable ToolContext ctx) { + + if (taskId == null || taskId.isBlank()) { + return errorJson("taskId 不能为空"); + } + String trimmedTaskId = taskId.trim(); + + AsyncTaskEntity entity = asyncTaskService.findEntityByTaskId(trimmedTaskId); + if (entity == null) { + return errorJson("Task not found: " + trimmedTaskId); + } + if (!"agent_delegate".equals(entity.getTaskType())) { + return errorJson("Task is not a delegate task: " + trimmedTaskId); + } + + // Attribution gate — registry is live-only, so the persistent + // request_json + created_by columns are the only authoritative + // sources. Both must match the calling context; otherwise this is a + // cross-user or cross-conversation lookup and must be denied even + // for an already-succeeded task (otherwise a stranger can read the + // result by guessing taskIds). + // + // Caveat on the user gate: when ChatOrigin.requesterId is empty, + // delegateAsync stamps the task with the literal sentinel "system" + // (mirrors the existing channel/cron-originated flow). All callers + // that share that sentinel — e.g. two cron jobs in the same + // workspace — therefore satisfy the user gate against each other. + // The conversation gate above still narrows it to "the same parent + // conversation as the spawn", which keeps the blast radius bounded; + // a follow-up that surfaces a stable per-channel / per-cron caller + // identity into ChatOrigin.requesterId would close this gap. + String taskParentConv; + try { + JsonNode req = entity.getRequestJson() == null + ? null + : objectMapper.readTree(entity.getRequestJson()); + taskParentConv = req == null ? "" : req.path("parentConversationId").asText(""); + } catch (Exception e) { + return errorJson("Failed to parse task payload: " + e.getMessage()); + } + String currentParentConv = resolveParentConversationId(); + ChatOrigin origin = ChatOrigin.from(ctx); + String currentUser = origin != null ? origin.requesterId() : null; + + if (taskParentConv.isEmpty() + || currentParentConv == null + || !taskParentConv.equals(currentParentConv)) { + return errorJson("Forbidden: task does not belong to current conversation"); + } + if (entity.getCreatedBy() == null || currentUser == null + || currentUser.isBlank() + || !entity.getCreatedBy().equals(currentUser)) { + return errorJson("Forbidden: task does not belong to current user"); + } + + String status = entity.getStatus(); + boolean isTerminal = "succeeded".equals(status) || "failed".equals(status); + if (Boolean.TRUE.equals(block) && !isTerminal) { + int waitSec = Math.min(TASK_OUTPUT_MAX_TIMEOUT_S, + Math.max(1, Optional.ofNullable(timeoutSeconds).orElse(TASK_OUTPUT_DEFAULT_TIMEOUT_S))); + long deadline = System.currentTimeMillis() + waitSec * 1000L; + while (System.currentTimeMillis() < deadline) { + try { + Thread.sleep(TASK_OUTPUT_POLL_INTERVAL_MS); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + break; + } + AsyncTaskEntity refreshed = asyncTaskService.findEntityByTaskId(trimmedTaskId); + if (refreshed == null) break; + entity = refreshed; + status = entity.getStatus(); + if ("succeeded".equals(status) || "failed".equals(status)) break; + } + } + + if (streamTracker.isRunning(currentParentConv)) { + streamTracker.broadcastObject(currentParentConv, "delegation_async_polled", Map.of( + "taskId", trimmedTaskId, + "status", status)); + } + + Map result = new LinkedHashMap<>(); + result.put("task_id", trimmedTaskId); + result.put("status", status); + switch (status == null ? "" : status) { + case "pending", "running" -> { + result.put("progress", entity.getProgress()); + result.put("hint", "Try again later or call task_output with block=true."); + } + case "succeeded" -> { + result.put("result", entity.getResultJson()); + result.put("duration_ms", durationMs(entity)); + } + case "failed" -> { + result.put("error", entity.getErrorMessage()); + result.put("duration_ms", durationMs(entity)); + } + default -> result.put("error", "Unknown status: " + status); + } + try { + return objectMapper.writeValueAsString(result); + } catch (Exception e) { + return errorJson("Failed to serialize response: " + e.getMessage()); + } + } + + /** Build a one-line JSON error envelope for tool returns. Kept distinct + * from {@link #truncate} / plain-text errors used by sync delegate paths + * so the model sees a consistent shape for async results. */ + private String errorJson(String message) { + try { + return objectMapper.writeValueAsString(Map.of( + "error", true, + "message", message != null ? message : "")); + } catch (Exception e) { + // Fallback — never throw from an error helper. + return "{\"error\":true,\"message\":\"" + (message == null ? "" : message.replace("\"", "\\\"")) + "\"}"; + } + } + + /** Walltime estimate using the create/update timestamps written by + * {@code AsyncTaskService}. Returns 0 when either timestamp is missing. */ + private static long durationMs(AsyncTaskEntity entity) { + if (entity == null || entity.getCreateTime() == null || entity.getUpdateTime() == null) return 0L; + return Duration.between(entity.getCreateTime(), entity.getUpdateTime()).toMillis(); + } + // ==================== Child agent execution (shared by single and parallel paths) ==================== /** diff --git a/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAgentToolDenyListTest.java b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAgentToolDenyListTest.java index 67e24718..355e71b6 100644 --- a/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAgentToolDenyListTest.java +++ b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAgentToolDenyListTest.java @@ -54,9 +54,10 @@ class DelegateAgentToolDenyListTest { ObjectMapper objectMapper = new ObjectMapper(); registry = new SubagentRegistry(); AuditEventService auditEventService = mock(AuditEventService.class); + vip.mate.task.AsyncTaskService asyncTaskService = mock(vip.mate.task.AsyncTaskService.class); tool = new DelegateAgentTool(agentService, agentMapper, streamTracker, conversationService, - objectMapper, registry, auditEventService); + objectMapper, registry, auditEventService, asyncTaskService); } @AfterEach diff --git a/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAsyncTaskOutputAttributionTest.java b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAsyncTaskOutputAttributionTest.java new file mode 100644 index 00000000..ae60d63f --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAsyncTaskOutputAttributionTest.java @@ -0,0 +1,203 @@ +package vip.mate.tool.builtin; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import org.springframework.ai.chat.model.ToolContext; +import vip.mate.agent.AgentService; +import vip.mate.agent.context.ChatOrigin; +import vip.mate.agent.delegation.SubagentRegistry; +import vip.mate.agent.repository.AgentMapper; +import vip.mate.audit.service.AuditEventService; +import vip.mate.channel.web.ChatStreamTracker; +import vip.mate.task.AsyncTaskService; +import vip.mate.task.model.AsyncTaskEntity; +import vip.mate.workspace.conversation.ConversationService; + +import java.time.LocalDateTime; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.when; + +/** + * Attribution gate for {@code taskOutput}. The persistent {@code mate_async_task} + * row carries both {@code created_by} (the requester) and a JSON blob whose + * {@code parentConversationId} field anchors the task to one conversation — + * any mismatch against the caller's current {@link ChatOrigin} must short-circuit + * to {@code Forbidden} before the row's body / result can leak. + *

+ * Three scenarios make up the threat model: + *

    + *
  • Cross-user — Alice's taskId is read by Bob in the same conversation.
    • + *
  • Cross-conversation — Alice reads her own taskId from a different + * conversation than the one that spawned it.
    • + *
  • Already-succeeded — same as above, but the row is terminal with a + * non-empty {@code result_json}; the failure mode here would leak the + * result body itself.
    • + *
+ */ +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class DelegateAsyncTaskOutputAttributionTest { + + @Mock private AgentService agentService; + @Mock private AgentMapper agentMapper; + @Mock private ChatStreamTracker streamTracker; + @Mock private ConversationService conversationService; + @Mock private SubagentRegistry subagentRegistry; + @Mock private AuditEventService auditEventService; + @Mock private AsyncTaskService asyncTaskService; + + private final ObjectMapper objectMapper = new ObjectMapper(); + private DelegateAgentTool tool; + + @BeforeEach + void setUp() { + tool = new DelegateAgentTool( + agentService, agentMapper, streamTracker, conversationService, + objectMapper, subagentRegistry, auditEventService, asyncTaskService); + } + + @AfterEach + void tearDown() { + ToolExecutionContext.clear(); + } + + @Test + @DisplayName("(a) Cross-user: task created by other-user → Forbidden, no result leaked") + void crossUserTaskIdForbidden() throws Exception { + // Task created by `other-user` in conv-shared. + AsyncTaskEntity entity = makeAsyncTask("tid-cross-user", "running", + "conv-shared", "other-user", null); + when(asyncTaskService.findEntityByTaskId("tid-cross-user")).thenReturn(entity); + + // Caller is user-1, sitting in conv-shared (so parentConv matches — + // only the user attribution should reject this). + ToolExecutionContext.set("conv-shared", "user-1"); + String result = tool.taskOutput("tid-cross-user", false, null, + makeCtx("user-1", "conv-shared")); + + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("current user"); + } + + @Test + @DisplayName("(b) Same-user different conversation: → Forbidden on conversation gate") + void sameUserDifferentConversationForbidden() throws Exception { + AsyncTaskEntity entity = makeAsyncTask("tid-cross-conv", "running", + "conv-A", "user-1", null); + when(asyncTaskService.findEntityByTaskId("tid-cross-conv")).thenReturn(entity); + + // user-1 is asking from conv-B; the task belongs to conv-A. + ToolExecutionContext.set("conv-B", "user-1"); + String result = tool.taskOutput("tid-cross-conv", false, null, + makeCtx("user-1", "conv-B")); + + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("current conversation"); + } + + @Test + @DisplayName("(c) Already-succeeded task accessed from wrong parent → Forbidden, no result body leaked") + void succeededTaskWrongParentForbidden() throws Exception { + // Succeeded row carries a non-empty result_json — exactly the body we + // must NOT echo back to a stranger guessing taskIds. + AsyncTaskEntity entity = makeAsyncTask("tid-done", "succeeded", + "conv-A", "user-1", "SECRET-ANSWER-PAYLOAD"); + when(asyncTaskService.findEntityByTaskId("tid-done")).thenReturn(entity); + + ToolExecutionContext.set("conv-B", "user-1"); + String result = tool.taskOutput("tid-done", false, null, + makeCtx("user-1", "conv-B")); + + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("current conversation"); + // Critical: the result body must not appear anywhere in the response. + assertThat(result).doesNotContain("SECRET-ANSWER-PAYLOAD"); + } + + @Test + @DisplayName("Legitimate caller (matching user + conversation) is allowed through") + void legitimateCallerAllowed() throws Exception { + AsyncTaskEntity entity = makeAsyncTask("tid-ok", "succeeded", + "conv-mine", "user-1", "valid result"); + when(asyncTaskService.findEntityByTaskId("tid-ok")).thenReturn(entity); + + ToolExecutionContext.set("conv-mine", "user-1"); + String result = tool.taskOutput("tid-ok", false, null, + makeCtx("user-1", "conv-mine")); + + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("status", "succeeded") + .containsEntry("result", "valid result"); + } + + @Test + @DisplayName("Empty parentConversationId in request_json → Forbidden (defensive)") + void emptyParentInPayloadForbidden() throws Exception { + AsyncTaskEntity entity = new AsyncTaskEntity(); + entity.setTaskId("tid-empty-parent"); + entity.setTaskType("agent_delegate"); + entity.setStatus("running"); + entity.setCreatedBy("user-1"); + // request_json with empty parentConversationId — should never happen + // in practice but the gate must still close. + entity.setRequestJson("{\"parentConversationId\":\"\",\"childConversationId\":\"child-x\"}"); + entity.setCreateTime(LocalDateTime.now()); + entity.setUpdateTime(LocalDateTime.now()); + when(asyncTaskService.findEntityByTaskId("tid-empty-parent")).thenReturn(entity); + + ToolExecutionContext.set("conv-X", "user-1"); + String result = tool.taskOutput("tid-empty-parent", false, null, + makeCtx("user-1", "conv-X")); + + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("current conversation"); + } + + // ---------- helpers ---------- + + private AsyncTaskEntity makeAsyncTask(String taskId, String status, String parentConv, + String createdBy, String resultJson) throws Exception { + AsyncTaskEntity e = new AsyncTaskEntity(); + e.setTaskId(taskId); + e.setTaskType("agent_delegate"); + e.setStatus(status); + e.setCreatedBy(createdBy); + e.setResultJson(resultJson); + e.setProgress("succeeded".equals(status) ? 100 : 50); + e.setCreateTime(LocalDateTime.now().minusSeconds(5)); + e.setUpdateTime(LocalDateTime.now()); + Map req = new LinkedHashMap<>(); + req.put("parentConversationId", parentConv); + req.put("childConversationId", "child-x"); + req.put("childAgentId", 10L); + req.put("task", "task"); + req.put("label", ""); + e.setRequestJson(objectMapper.writeValueAsString(req)); + return e; + } + + private ToolContext makeCtx(String requester, String conversationId) { + ChatOrigin origin = new ChatOrigin( + 1L, conversationId, requester, null, null, null, null); + Map map = new HashMap<>(); + map.put(ChatOrigin.CTX_KEY, origin); + return new ToolContext(map); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAsyncToolTest.java b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAsyncToolTest.java new file mode 100644 index 00000000..719928a7 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DelegateAsyncToolTest.java @@ -0,0 +1,347 @@ +package vip.mate.tool.builtin; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import org.springframework.ai.chat.model.ToolContext; +import vip.mate.agent.AgentService; +import vip.mate.agent.context.ChatOrigin; +import vip.mate.agent.delegation.SubagentRegistry; +import vip.mate.agent.model.AgentEntity; +import vip.mate.agent.repository.AgentMapper; +import vip.mate.audit.service.AuditEventService; +import vip.mate.channel.web.ChatStreamTracker; +import vip.mate.task.AsyncTaskService; +import vip.mate.task.model.AsyncTaskEntity; +import vip.mate.workspace.conversation.ConversationService; + +import java.time.LocalDateTime; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeast; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Behavioral contract for the two async-delegation tools — + * {@code delegateAsync} (spawn returns task_id immediately) and + * {@code taskOutput} (status / result retrieval). + *

+ * AsyncTaskService is mocked, so the Callable submitted by delegateAsync is + * never invoked here: the inner execution path is covered by + * {@code AsyncTaskServiceOneShotTest}. What this suite locks down is + * the synchronous shell — argument validation, depth / spawn-pause guards, + * cap-overflow degradation, JSON shape, and the SSE spawn-event side effect. + */ +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class DelegateAsyncToolTest { + + @Mock private AgentService agentService; + @Mock private AgentMapper agentMapper; + @Mock private ChatStreamTracker streamTracker; + @Mock private ConversationService conversationService; + @Mock private SubagentRegistry subagentRegistry; + @Mock private AuditEventService auditEventService; + @Mock private AsyncTaskService asyncTaskService; + + private final ObjectMapper objectMapper = new ObjectMapper(); + private DelegateAgentTool tool; + + @BeforeEach + void setUp() { + tool = new DelegateAgentTool( + agentService, agentMapper, streamTracker, conversationService, + objectMapper, subagentRegistry, auditEventService, asyncTaskService); + // resolveParentConversationId reads from ToolExecutionContext first; + // seed it so the async delegation has a parent to attach the task to. + ToolExecutionContext.set("parent-conv-1", "user-1"); + } + + @AfterEach + void tearDown() { + ToolExecutionContext.clear(); + while (DelegationContext.currentDepth() > 0) { + DelegationContext.exit(); + } + } + + // ---------- delegateAsync ---------- + + @Test + @DisplayName("delegateAsync returns task_id, child_conversation_id, status=running synchronously") + @SuppressWarnings("unchecked") + void delegateAsyncReturnsTaskIdImmediately() throws Exception { + AgentEntity target = makeAgent(10L, "Researcher"); + when(agentMapper.selectOne(any())).thenReturn(target); + when(subagentRegistry.isSpawnPaused("parent-conv-1")).thenReturn(false); + when(subagentRegistry.register(anyString(), anyString(), anyLong(), anyString(), any())) + .thenReturn("sa-1"); + + AsyncTaskEntity entity = new AsyncTaskEntity(); + entity.setTaskId("tid-123"); + when(asyncTaskService.submitOneShot( + eq("agent_delegate"), eq("parent-conv-1"), any(), anyString(), eq("user-1"), any())) + .thenReturn(entity); + when(streamTracker.isRunning("parent-conv-1")).thenReturn(true); + + String result = tool.delegateAsync("Researcher", "Go research things", "label-x", makeCtx("user-1", "parent-conv-1")); + + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("task_id", "tid-123") + .containsEntry("status", "running") + .containsEntry("agent_name", "Researcher") + .containsEntry("label", "label-x"); + assertThat((String) parsed.get("child_conversation_id")).startsWith("child-"); + assertThat((String) parsed.get("hint")).contains("task_output"); + + // The spawn SSE event reaches the parent's stream. + verify(streamTracker).broadcastObject(eq("parent-conv-1"), + eq("delegation_async_spawned"), any(Map.class)); + } + + @Test + @DisplayName("delegateAsync passes a request_json payload carrying parent + child + agentId + label") + void delegateAsyncRequestJsonShape() throws Exception { + AgentEntity target = makeAgent(10L, "Researcher"); + when(agentMapper.selectOne(any())).thenReturn(target); + when(subagentRegistry.register(anyString(), anyString(), anyLong(), anyString(), any())) + .thenReturn("sa-2"); + AsyncTaskEntity entity = new AsyncTaskEntity(); + entity.setTaskId("tid-200"); + when(asyncTaskService.submitOneShot(anyString(), anyString(), any(), anyString(), anyString(), any())) + .thenReturn(entity); + + tool.delegateAsync("Researcher", "task body", "myLabel", makeCtx("user-1", "parent-conv-1")); + + org.mockito.ArgumentCaptor jsonCaptor = org.mockito.ArgumentCaptor.forClass(String.class); + verify(asyncTaskService).submitOneShot( + eq("agent_delegate"), eq("parent-conv-1"), any(), + jsonCaptor.capture(), eq("user-1"), any()); + Map payload = objectMapper.readValue(jsonCaptor.getValue(), new TypeReference<>() {}); + assertThat(payload).containsEntry("parentConversationId", "parent-conv-1") + .containsEntry("label", "myLabel") + .containsEntry("task", "task body"); + assertThat(payload.get("childConversationId")).asString().startsWith("child-"); + assertThat(((Number) payload.get("childAgentId")).longValue()).isEqualTo(10L); + } + + @Test + @DisplayName("Concurrency-cap (IllegalStateException) → error JSON + registry unregistered") + void delegateAsyncConcurrencyCap() throws Exception { + AgentEntity target = makeAgent(10L, "Researcher"); + when(agentMapper.selectOne(any())).thenReturn(target); + when(subagentRegistry.register(anyString(), anyString(), anyLong(), anyString(), any())) + .thenReturn("sa-cap"); + when(asyncTaskService.submitOneShot(anyString(), anyString(), any(), anyString(), anyString(), any())) + .thenThrow(new IllegalStateException("已达到最大并行任务数（3），请等待现有任务完成")); + + String result = tool.delegateAsync("Researcher", "task", null, makeCtx("user-1", "parent-conv-1")); + + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("最大并行任务数"); + + // Registry entry MUST be released so it doesn't dangle through the cap. + verify(subagentRegistry).unregister("sa-cap"); + // No spawn event broadcast for a failed spawn. + verify(streamTracker, never()).broadcastObject(anyString(), + eq("delegation_async_spawned"), any()); + } + + @Test + @DisplayName("Missing agentName / task → error JSON without touching downstream services") + void delegateAsyncMissingArgs() throws Exception { + String r1 = tool.delegateAsync("", "task", null, makeCtx("user-1", "parent-conv-1")); + String r2 = tool.delegateAsync("X", " ", null, makeCtx("user-1", "parent-conv-1")); + for (String r : new String[]{r1, r2}) { + Map parsed = objectMapper.readValue(r, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + } + verify(asyncTaskService, never()).submitOneShot(any(), any(), any(), any(), any(), any()); + verify(subagentRegistry, never()).register(any(), any(), any(), any(), any()); + } + + @Test + @DisplayName("Agent not found → error JSON") + void delegateAsyncAgentNotFound() throws Exception { + when(agentMapper.selectOne(any())).thenReturn(null); + String result = tool.delegateAsync("Ghost", "task", null, makeCtx("user-1", "parent-conv-1")); + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("Ghost"); + verify(asyncTaskService, never()).submitOneShot(any(), any(), any(), any(), any(), any()); + } + + @Test + @DisplayName("Spawn-pause active → error JSON, no task submitted, no registry entry") + void delegateAsyncSpawnPause() throws Exception { + AgentEntity target = makeAgent(10L, "Researcher"); + when(agentMapper.selectOne(any())).thenReturn(target); + when(subagentRegistry.isSpawnPaused("parent-conv-1")).thenReturn(true); + + String result = tool.delegateAsync("Researcher", "task", null, makeCtx("user-1", "parent-conv-1")); + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("paused"); + verify(asyncTaskService, never()).submitOneShot(any(), any(), any(), any(), any(), any()); + verify(subagentRegistry, never()).register(any(), any(), any(), any(), any()); + } + + @Test + @DisplayName("Depth limit reached → error JSON") + void delegateAsyncDepthLimit() throws Exception { + // Push depth to MAX (3) so currentDepth >= MAX_DELEGATION_DEPTH. + for (int i = 0; i < 3; i++) { + DelegationContext.enter("parent-conv-1", java.util.Set.of()); + } + String result = tool.delegateAsync("Researcher", "task", null, makeCtx("user-1", "parent-conv-1")); + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("depth"); + verify(asyncTaskService, never()).submitOneShot(any(), any(), any(), any(), any(), any()); + } + + // ---------- taskOutput ---------- + + @Test + @DisplayName("taskOutput on running task with block=false returns status=running") + void taskOutputRunning() throws Exception { + AsyncTaskEntity entity = makeAsyncTask("tid-run", "running", "parent-conv-1", "user-1", null); + when(asyncTaskService.findEntityByTaskId("tid-run")).thenReturn(entity); + + String result = tool.taskOutput("tid-run", false, null, makeCtx("user-1", "parent-conv-1")); + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("status", "running") + .containsEntry("task_id", "tid-run"); + assertThat((String) parsed.get("hint")).contains("Try again"); + } + + @Test + @DisplayName("taskOutput on succeeded task returns result + duration_ms") + void taskOutputSucceeded() throws Exception { + AsyncTaskEntity entity = makeAsyncTask("tid-ok", "succeeded", "parent-conv-1", "user-1", "child final answer"); + when(asyncTaskService.findEntityByTaskId("tid-ok")).thenReturn(entity); + + String result = tool.taskOutput("tid-ok", null, null, makeCtx("user-1", "parent-conv-1")); + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("status", "succeeded") + .containsEntry("result", "child final answer"); + assertThat(((Number) parsed.get("duration_ms")).longValue()).isGreaterThanOrEqualTo(0L); + } + + @Test + @DisplayName("taskOutput on failed task returns error message") + void taskOutputFailed() throws Exception { + AsyncTaskEntity entity = makeAsyncTask("tid-fail", "failed", "parent-conv-1", "user-1", null); + entity.setErrorMessage("agent boom"); + when(asyncTaskService.findEntityByTaskId("tid-fail")).thenReturn(entity); + + String result = tool.taskOutput("tid-fail", false, null, makeCtx("user-1", "parent-conv-1")); + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("status", "failed") + .containsEntry("error", "agent boom"); + } + + @Test + @DisplayName("Unknown taskId → error JSON without touching parent SSE") + void taskOutputNotFound() throws Exception { + when(asyncTaskService.findEntityByTaskId("tid-missing")).thenReturn(null); + String result = tool.taskOutput("tid-missing", false, null, makeCtx("user-1", "parent-conv-1")); + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("Task not found"); + verify(streamTracker, never()).broadcastObject(any(), + eq("delegation_async_polled"), any()); + } + + @Test + @DisplayName("Non-agent_delegate taskType (e.g. video_generation) → error JSON") + void taskOutputWrongTaskType() throws Exception { + AsyncTaskEntity entity = new AsyncTaskEntity(); + entity.setTaskId("tid-vid"); + entity.setTaskType("video_generation"); + entity.setStatus("running"); + when(asyncTaskService.findEntityByTaskId("tid-vid")).thenReturn(entity); + + String result = tool.taskOutput("tid-vid", false, null, makeCtx("user-1", "parent-conv-1")); + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("error", true); + assertThat((String) parsed.get("message")).contains("not a delegate task"); + } + + @Test + @DisplayName("block=true on a task that stays running for the full timeout returns status=running") + void taskOutputBlockTimeout() throws Exception { + AsyncTaskEntity entity = makeAsyncTask("tid-block", "running", "parent-conv-1", "user-1", null); + when(asyncTaskService.findEntityByTaskId("tid-block")).thenReturn(entity); + + long start = System.currentTimeMillis(); + // 1s budget → poll loop runs ~2 iterations of 500ms before deadline. + String result = tool.taskOutput("tid-block", true, 1, makeCtx("user-1", "parent-conv-1")); + long elapsed = System.currentTimeMillis() - start; + + Map parsed = objectMapper.readValue(result, new TypeReference<>() {}); + assertThat(parsed).containsEntry("status", "running"); + // Real-time guard: ≥ ~900 ms (loop ran) but well under the 120 s cap. + assertThat(elapsed).isBetween(900L, 5_000L); + // Initial read + at least one poll iteration. + verify(asyncTaskService, atLeast(2)).findEntityByTaskId("tid-block"); + } + + // ---------- helpers ---------- + + private static AgentEntity makeAgent(Long id, String name) { + AgentEntity a = new AgentEntity(); + a.setId(id); + a.setName(name); + a.setEnabled(true); + a.setWorkspaceId(1L); + return a; + } + + private AsyncTaskEntity makeAsyncTask(String taskId, String status, String parentConv, + String createdBy, String resultJson) throws Exception { + AsyncTaskEntity e = new AsyncTaskEntity(); + e.setTaskId(taskId); + e.setTaskType("agent_delegate"); + e.setStatus(status); + e.setCreatedBy(createdBy); + e.setResultJson(resultJson); + e.setProgress("running".equals(status) ? 50 : ("succeeded".equals(status) ? 100 : 0)); + e.setCreateTime(LocalDateTime.now().minusSeconds(5)); + e.setUpdateTime(LocalDateTime.now()); + Map req = new LinkedHashMap<>(); + req.put("parentConversationId", parentConv); + req.put("childConversationId", "child-x"); + req.put("childAgentId", 10L); + req.put("task", "task"); + req.put("label", ""); + e.setRequestJson(objectMapper.writeValueAsString(req)); + return e; + } + + private ToolContext makeCtx(String requester, String conversationId) { + ChatOrigin origin = new ChatOrigin( + 1L, conversationId, requester, null, null, null, null); + Map map = new HashMap<>(); + map.put(ChatOrigin.CTX_KEY, origin); + return new ToolContext(map); + } +}