diff --git a/mateclaw-server/src/main/java/vip/mate/task/AsyncTaskService.java b/mateclaw-server/src/main/java/vip/mate/task/AsyncTaskService.java index 54647533..60e383fe 100644 --- a/mateclaw-server/src/main/java/vip/mate/task/AsyncTaskService.java +++ b/mateclaw-server/src/main/java/vip/mate/task/AsyncTaskService.java @@ -133,6 +133,130 @@ public class AsyncTaskService implements ApplicationRunner { return entity; } + // ==================== One-shot Callable submission ==================== + + /** + * Submit a one-shot {@link Callable} that runs on this service's + * {@code pollExecutor} and persists its outcome through the standard + * {@code mate_async_task} lifecycle. + *

+ * This is the local-work counterpart to {@link #startPolling}: instead of + * polling an external provider, the worker runs {@code work.call()} in + * the shared executor, writes the return value to {@code resultJson} (or + * the exception message to {@code errorMessage}), and registers itself in + * the same {@code activePolls} / {@code pollTaskToConv} bookkeeping so + * {@link #onConversationDeleted} can cancel both kinds uniformly. + *

+ * Race closure: the worker is scheduled at 0 ms but blocks on an internal + * {@code CountDownLatch} until the calling thread has finished + * registering both bookkeeping entries. Without this, the executor could + * dequeue the worker, run its {@code finally} cleanup, and return — all + * before {@code activePolls.put} runs on the calling thread — leaving a + * ghost entry no later event drains. + *

+ * Cancellation: while {@code work.call()} runs the worker has no + * cooperative cancel signal beyond {@link #isConversationCanceled}; it + * checks once before invoking the body and once after, so a parent + * conversation deleted mid-run never lands as {@code succeeded}. + * {@link #onConversationDeleted} additionally writes {@code failed} + * synchronously for any non-terminal {@code agent_delegate} task so the + * DB row never lingers in {@code running} after parent deletion. + * + * @param taskType Discriminator written to {@code task_type} + * (e.g. {@code "agent_delegate"}). Listeners + * and the conversation-deleted DB write-back gate + * on this value. + * @param conversationId Parent conversation ID. Written to + * {@code conversation_id} so deleting the parent + * conversation cascade-cancels the worker. Any + * child / detached identifiers belong in + * {@code requestJson}, not here. + * @param messageId Optional parent message ID. + * @param requestJson Caller-serialized request payload. + * @param createdBy Audit attribution; counts toward + * {@code MAX_ACTIVE_TASKS_PER_USER}. + * @param work Body whose return value is persisted as + * {@code resultJson}. A thrown exception lands as + * {@code status=failed} with the exception + * message recorded. + * @return the created task entity (status = pending at return time). + */ + public AsyncTaskEntity submitOneShot(String taskType, String conversationId, + Long messageId, String requestJson, + String createdBy, Callable work) { + AsyncTaskEntity entity = createTask(taskType, conversationId, messageId, + "internal", null, requestJson, createdBy); + final String taskId = entity.getTaskId(); + updateStatus(taskId, "running", 0, null, null); + + // schedule(0)-vs-put race closure: see method javadoc. + CountDownLatch enrolled = new CountDownLatch(1); + ScheduledFuture future = pollExecutor.schedule(() -> { + try { + try { + enrolled.await(); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + updateStatus(taskId, "failed", null, null, "worker interrupted before start"); + return; + } + // Pre-call cancel check: parent conversation may have been + // deleted between submitOneShot returning and the worker + // being dequeued. + if (isConversationCanceled(conversationId)) { + updateStatus(taskId, "failed", null, null, "conversation deleted before start"); + return; + } + String result; + try { + result = work.call(); + } catch (Exception e) { + log.warn("[AsyncTask] One-shot task {} failed: {}", taskId, e.getMessage()); + updateStatus(taskId, "failed", null, null, e.getMessage()); + return; + } + // Post-call cancel check: parent may have been deleted while + // work was running; avoid resurrecting a succeeded row for a + // conversation whose DB cascade already wiped it. + if (isConversationCanceled(conversationId)) { + updateStatus(taskId, "failed", null, null, "conversation deleted during execution"); + } else { + updateStatus(taskId, "succeeded", 100, result, null); + } + } finally { + activePolls.remove(taskId); + pollTaskToConv.remove(taskId); + // Generic terminal event so SSE listeners (parent + // conversation of an async delegation, UI badges, …) can + // react without polling the DB. Re-fetch so the event + // payload reflects the row we just wrote. Wrapped in a + // try-catch because a broadcast failure on a stale + // conversation stream must not mask the task outcome. + try { + AsyncTaskEntity finalEntity = findEntityByTaskId(taskId); + if (finalEntity != null) { + boolean success = "succeeded".equals(finalEntity.getStatus()); + broadcastTaskEventWithData(finalEntity, "async_task_completed", + success, java.util.Map.of(), + success ? null : finalEntity.getErrorMessage()); + } + } catch (Exception broadcastErr) { + log.debug("[AsyncTask] Completion broadcast failed for task {}: {}", + taskId, broadcastErr.getMessage()); + } + } + }, 0, TimeUnit.MILLISECONDS); + + activePolls.put(taskId, future); + if (conversationId != null) { + pollTaskToConv.put(taskId, conversationId); + } + enrolled.countDown(); + log.info("[AsyncTask] Submitted one-shot task {} (type={}, conv={})", + taskId, taskType, conversationId); + return entity; + } + // ==================== 轮询管理 ==================== /** @@ -145,7 +269,7 @@ public class AsyncTaskService implements ApplicationRunner { public void startPolling(String taskId, Function statusChecker, BiConsumer onComplete) { - AsyncTaskEntity task = findByTaskId(taskId); + AsyncTaskEntity task = findEntityByTaskId(taskId); if (task == null) { log.warn("[AsyncTask] Cannot start polling: task {} not found", taskId); return; @@ -190,7 +314,7 @@ public class AsyncTaskService implements ApplicationRunner { updateStatus(taskId, "failed", null, null, result.errorMessage()); } // 刷新任务实体 - AsyncTaskEntity freshTask = findByTaskId(taskId); + AsyncTaskEntity freshTask = findEntityByTaskId(taskId); onComplete.accept(freshTask, result); } } catch (Exception e) { @@ -251,7 +375,22 @@ public class AsyncTaskService implements ApplicationRunner { int cancelled = 0; for (Map.Entry entry : pollTaskToConv.entrySet()) { if (convId.equals(entry.getValue())) { - cancelPolling(entry.getKey()); + String taskId = entry.getKey(); + cancelPolling(taskId); + // One-shot tasks (taskType "agent_delegate") have no separate + // poll loop to observe the cancel and write the terminal row: + // cancelPolling only nukes the Future. Without this explicit + // write-back the DB row stays "running" forever. Polling + // tasks (video / image / ...) keep their original behavior — + // their own poll completion or startup-recovery path is what + // writes the terminal status. + AsyncTaskEntity t = findEntityByTaskId(taskId); + if (t != null + && "agent_delegate".equals(t.getTaskType()) + && !"succeeded".equals(t.getStatus()) + && !"failed".equals(t.getStatus())) { + updateStatus(taskId, "failed", null, null, "conversation deleted"); + } cancelled++; } } @@ -290,7 +429,7 @@ public class AsyncTaskService implements ApplicationRunner { // ==================== 查询 ==================== public AsyncTaskInfo getTaskInfo(String taskId) { - AsyncTaskEntity entity = findByTaskId(taskId); + AsyncTaskEntity entity = findEntityByTaskId(taskId); if (entity == null) { return null; } @@ -307,7 +446,12 @@ public class AsyncTaskService implements ApplicationRunner { return entities.stream().map(this::toInfo).toList(); } - private AsyncTaskEntity findByTaskId(String taskId) { + /** Returns the persisted task entity by its public {@code taskId}, or + * {@code null} if no row matches. Promoted from private to public so the + * conversation-deleted listener (and overrides in tests) can resolve a + * task's current state without going through the read-model + * {@link #getTaskInfo}. */ + public AsyncTaskEntity findEntityByTaskId(String taskId) { return asyncTaskMapper.selectOne( new LambdaQueryWrapper() .eq(AsyncTaskEntity::getTaskId, taskId) diff --git a/mateclaw-server/src/test/java/vip/mate/task/AsyncTaskServiceOneShotTest.java b/mateclaw-server/src/test/java/vip/mate/task/AsyncTaskServiceOneShotTest.java new file mode 100644 index 00000000..f6e18cd5 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/task/AsyncTaskServiceOneShotTest.java @@ -0,0 +1,347 @@ +package vip.mate.task; + +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.springframework.test.util.ReflectionTestUtils; +import vip.mate.channel.web.ChatStreamTracker; +import vip.mate.task.model.AsyncTaskEntity; +import vip.mate.task.repository.AsyncTaskMapper; +import vip.mate.workspace.conversation.event.ConversationDeletedEvent; + +import java.util.UUID; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BooleanSupplier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; + +/** + * Contract for {@link AsyncTaskService#submitOneShot} — the one-shot Callable + * entry point that lets a caller hand off arbitrary work to {@code pollExecutor} + * and have the standard {@code mate_async_task} lifecycle (running → terminal + * + automatic cancel-on-parent-conversation-deleted + bookkeeping cleanup) take + * over. + *

+ * The fixture installs an anonymous subclass that overrides {@code createTask}, + * {@code updateStatus} and {@code findEntityByTaskId} with in-memory equivalents. + * Reason: the production methods round-trip through a MyBatis-Plus + * LambdaUpdateWrapper that obscures the {@code (status, progress, resultJson, + * errorMessage)} tuple under MPGENVAL placeholders. Asserting on a + * test-controlled state map is faster, less brittle, and lets the suite focus + * on what {@code submitOneShot} actually does — schedule the worker, register + * bookkeeping, observe cancellation, and clean up. + * + *

Covered paths

+ *
    + *
  1. Success — Callable returns, status lands on succeeded with resultJson.
    2. + *
  2. Exception — Callable throws, status lands on failed with errorMessage.
    3. + *
  3. Conversation deletion mid-run — listener writes failed + worker's + * second cancel-check also writes failed; both messages match.
    4. + *
  4. schedule/put race stress — 200 zero-cost tasks all succeed and drain + * both bookkeeping maps to empty (catches any ghost entry).
    5. + *
  5. Latch ordering — worker observes itself enrolled in both maps the + * moment Callable.call() begins (the latch invariant).
    6. + *
+ */ +class AsyncTaskServiceOneShotTest { + + private AsyncTaskMapper mapper; + private ChatStreamTracker tracker; + + /** Test-side per-taskId snapshot — populated by the fixture's overrides + * of createTask + updateStatus instead of going through MyBatis-Plus. */ + private static final class TaskState { + String taskId; + String taskType; + String conversationId; + String status; + Integer progress; + String resultJson; + String errorMessage; + } + + private final ConcurrentMap states = new ConcurrentHashMap<>(); + private AsyncTaskService service; + + @BeforeEach + void setUp() { + mapper = mock(AsyncTaskMapper.class); + tracker = mock(ChatStreamTracker.class); + states.clear(); + service = new AsyncTaskService(mapper, tracker) { + @Override + public AsyncTaskEntity createTask(String taskType, String conversationId, Long messageId, + String providerName, String providerTaskId, + String requestJson, String createdBy) { + String taskId = UUID.randomUUID().toString().replace("-", "").substring(0, 16); + AsyncTaskEntity e = new AsyncTaskEntity(); + e.setTaskId(taskId); + e.setTaskType(taskType); + e.setStatus("pending"); + e.setConversationId(conversationId); + e.setMessageId(messageId); + e.setProviderName(providerName); + e.setProviderTaskId(providerTaskId); + e.setRequestJson(requestJson); + e.setCreatedBy(createdBy); + TaskState s = new TaskState(); + s.taskId = taskId; + s.taskType = taskType; + s.conversationId = conversationId; + s.status = "pending"; + states.put(taskId, s); + return e; + } + + @Override + public void updateStatus(String taskId, String status, Integer progress, + String resultJson, String errorMessage) { + states.compute(taskId, (k, old) -> { + TaskState s = old != null ? old : new TaskState(); + s.taskId = taskId; + s.status = status; + if (progress != null) s.progress = progress; + if (resultJson != null) s.resultJson = resultJson; + if (errorMessage != null) s.errorMessage = errorMessage; + return s; + }); + } + + @Override + public AsyncTaskEntity findEntityByTaskId(String taskId) { + TaskState s = states.get(taskId); + if (s == null) return null; + AsyncTaskEntity e = new AsyncTaskEntity(); + e.setTaskId(s.taskId); + e.setTaskType(s.taskType); + e.setStatus(s.status); + e.setConversationId(s.conversationId); + return e; + } + }; + } + + @AfterEach + void tearDown() { + service.shutdown(); + } + + @Test + @DisplayName("Success path: status = succeeded, resultJson captures Callable return, maps drain") + void successPath() throws Exception { + AsyncTaskEntity entity = service.submitOneShot( + "agent_delegate", "conv-success", null, "{}", "user-1", + () -> "ok"); + + awaitDone(entity.getTaskId(), 5_000); + + TaskState s = states.get(entity.getTaskId()); + assertThat(s.status).isEqualTo("succeeded"); + assertThat(s.resultJson).isEqualTo("ok"); + assertThat(s.progress).isEqualTo(100); + assertActiveMapsEmpty(); + } + + @Test + @DisplayName("Exception path: status = failed, errorMessage carries Callable's message") + void exceptionPath() throws Exception { + AsyncTaskEntity entity = service.submitOneShot( + "agent_delegate", "conv-fail", null, "{}", "user-1", + () -> { throw new RuntimeException("boom-msg"); }); + + awaitDone(entity.getTaskId(), 5_000); + + TaskState s = states.get(entity.getTaskId()); + assertThat(s.status).isEqualTo("failed"); + assertThat(s.errorMessage).isNotNull().contains("boom-msg"); + assertActiveMapsEmpty(); + } + + @Test + @DisplayName("Conversation deleted while running: terminal status is failed with deletion message") + void cancelPathViaConversationDeleted() throws Exception { + CountDownLatch workerStarted = new CountDownLatch(1); + CountDownLatch workerCanProceed = new CountDownLatch(1); + + String convId = "conv-cancel"; + AsyncTaskEntity entity = service.submitOneShot( + "agent_delegate", convId, null, "{}", "user-1", + () -> { + workerStarted.countDown(); + // Hold the worker inside work.call() until the test fires + // the deletion event. Future.cancel(false) — what the + // listener calls — does NOT interrupt, so this await won't + // unblock until the test's explicit countDown below. + workerCanProceed.await(5, TimeUnit.SECONDS); + return "should-not-be-applied"; + }); + + assertThat(workerStarted.await(5, TimeUnit.SECONDS)).isTrue(); + + // Listener runs synchronously on the test thread: it cancels the + // future, looks up the (still-running) entity, sees taskType + // agent_delegate and writes failed + "conversation deleted". This + // closes the contract gap that bare cancelPolling left open before + // this PR (DB row would otherwise stay running forever). + service.onConversationDeleted(new ConversationDeletedEvent(convId)); + + // Release the worker so it observes isConversationCanceled = true at + // the post-call cancel check and writes the during-execution variant + // (also matches "conversation deleted"). + workerCanProceed.countDown(); + awaitDone(entity.getTaskId(), 5_000); + + TaskState s = states.get(entity.getTaskId()); + assertThat(s.status).isEqualTo("failed"); + assertThat(s.errorMessage).isNotNull().contains("conversation deleted"); + assertActiveMapsEmpty(); + } + + @Test + @DisplayName("Terminal status broadcasts async_task_completed on the parent conversation") + void broadcastsCompletionEvent() throws Exception { + AsyncTaskEntity successEntity = service.submitOneShot( + "agent_delegate", "conv-broadcast-ok", null, "{}", "user-1", + () -> "ok"); + AsyncTaskEntity failEntity = service.submitOneShot( + "agent_delegate", "conv-broadcast-fail", null, "{}", "user-1", + () -> { throw new RuntimeException("boom"); }); + + awaitDone(successEntity.getTaskId(), 5_000); + awaitDone(failEntity.getTaskId(), 5_000); + + // Success path → event with success=true, no errorMessage. + verify(tracker, timeout(2_000)).broadcastObject( + eq("conv-broadcast-ok"), eq("async_task_completed"), any()); + // Failure path → event with success=false, errorMessage carries + // the Callable's exception message. Broadcast routes to the parent + // conversation_id stored on the entity, which IS the parent for + // agent-delegate one-shots per AsyncTaskService.submitOneShot. + verify(tracker, timeout(2_000)).broadcastObject( + eq("conv-broadcast-fail"), eq("async_task_completed"), any()); + } + + @Test + @DisplayName("schedule/put race stress: 200 zero-cost tasks all succeed and bookkeeping drains") + void scheduleAndPutRaceStress() throws Exception { + int iterations = 200; + AsyncTaskEntity[] entities = new AsyncTaskEntity[iterations]; + for (int i = 0; i < iterations; i++) { + entities[i] = service.submitOneShot( + "agent_delegate", "conv-race-" + i, null, "{}", "user-1", + () -> "ok"); + } + for (AsyncTaskEntity e : entities) { + awaitDone(e.getTaskId(), 15_000); + } + // All terminal statuses must be succeeded. + for (AsyncTaskEntity e : entities) { + assertThat(states.get(e.getTaskId()).status) + .as("task %s succeeded", e.getTaskId()) + .isEqualTo("succeeded"); + } + // Belt-and-suspenders: any ghost (future.isDone() but key still in + // map) would keep one or both of these non-empty. + awaitMapsDrained(5_000); + assertActiveMapsEmpty(); + } + + @Test + @DisplayName("Latch ordering: Callable.call() observes its taskId enrolled in both maps") + void enrolledLatchOrdering() throws Exception { + ConcurrentHashMap activePolls = getInternalMap("activePolls"); + ConcurrentHashMap pollTaskToConv = getInternalMap("pollTaskToConv"); + + AtomicReference taskIdRef = new AtomicReference<>(); + AtomicBoolean observedActive = new AtomicBoolean(); + AtomicBoolean observedConvLink = new AtomicBoolean(); + CountDownLatch probed = new CountDownLatch(1); + + // The Callable acts as a probe: if the latch invariant holds, both + // bookkeeping maps already contain this taskId by the time the body + // runs (calling thread put before countDown, worker awaited + // countDown before reaching here). If the latch were removed, the + // worker could race ahead, observe empty maps, and finish — then the + // calling thread's put would leak a ghost entry. + Callable work = () -> { + String tid = taskIdRef.get(); + observedActive.set(tid != null && activePolls.containsKey(tid)); + observedConvLink.set(tid != null && pollTaskToConv.containsKey(tid)); + probed.countDown(); + return "ok"; + }; + + AsyncTaskEntity entity = service.submitOneShot( + "agent_delegate", "conv-latch", null, "{}", "user-1", work); + taskIdRef.set(entity.getTaskId()); + + assertThat(probed.await(5, TimeUnit.SECONDS)) + .as("probe must run within timeout") + .isTrue(); + awaitDone(entity.getTaskId(), 5_000); + + assertThat(observedActive) + .as("activePolls must contain taskId when work.call() begins") + .isTrue(); + assertThat(observedConvLink) + .as("pollTaskToConv must contain taskId when work.call() begins") + .isTrue(); + assertActiveMapsEmpty(); + } + + // ---------- helpers ---------- + + /** Waits until the task has reached a terminal status AND the worker's + * finally cleanup has cleared this taskId from both bookkeeping maps. */ + private void awaitDone(String taskId, long timeoutMs) throws InterruptedException { + awaitUntil(() -> { + TaskState s = states.get(taskId); + if (s == null) return false; + if (!"succeeded".equals(s.status) && !"failed".equals(s.status)) return false; + return !getInternalMap("activePolls").containsKey(taskId) + && !getInternalMap("pollTaskToConv").containsKey(taskId); + }, timeoutMs, "task " + taskId + " never reached cleaned-up terminal state"); + } + + private void awaitMapsDrained(long timeoutMs) throws InterruptedException { + awaitUntil(() -> getInternalMap("activePolls").isEmpty() + && getInternalMap("pollTaskToConv").isEmpty(), + timeoutMs, "activePolls / pollTaskToConv never drained"); + } + + private void awaitUntil(BooleanSupplier cond, long timeoutMs, String message) + throws InterruptedException { + long deadline = System.currentTimeMillis() + timeoutMs; + while (System.currentTimeMillis() < deadline) { + if (cond.getAsBoolean()) return; + Thread.sleep(20); + } + throw new AssertionError(message + " (waited " + timeoutMs + "ms)"); + } + + private void assertActiveMapsEmpty() { + assertThat(getInternalMap("activePolls")) + .as("activePolls should be empty after all workers finish") + .isEmpty(); + assertThat(getInternalMap("pollTaskToConv")) + .as("pollTaskToConv should be empty after all workers finish") + .isEmpty(); + } + + @SuppressWarnings("unchecked") + private ConcurrentHashMap getInternalMap(String name) { + return (ConcurrentHashMap) ReflectionTestUtils.getField(service, name); + } +}