mateclaw/mateclaw-server/src/test/java/vip/mate/goal/service/GoalSegmentRunnerTest.java

419 lines
24 KiB
Java

package vip.mate.goal.service;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import vip.mate.agent.AgentService;
import vip.mate.agent.context.GoalContinuationContext;
import vip.mate.agent.model.AgentEntity;
import vip.mate.agent.runtime.ConversationTurnGate;
import vip.mate.approval.ApprovalWorkflowService;
import vip.mate.channel.web.ChatStreamTracker;
import vip.mate.channel.web.ConversationInputQueueStore;
import vip.mate.goal.model.GoalEntity;
import vip.mate.goal.model.SegmentOutcome;
import vip.mate.workspace.conversation.ConversationService;
import vip.mate.workspace.conversation.model.ConversationEntity;
import vip.mate.workspace.conversation.model.MessageContentPart;
import vip.mate.workspace.conversation.model.MessageEntity;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;
class GoalSegmentRunnerTest {
AgentService agents=mock(AgentService.class);
ConversationService conversations=mock(ConversationService.class);
ApprovalWorkflowService approvals=mock(ApprovalWorkflowService.class);
ChatStreamTracker streams=new ChatStreamTracker(new ObjectMapper());
ConversationInputQueueStore inputQueue=mock(ConversationInputQueueStore.class);
ConcurrentLinkedQueue<ConversationInputQueueStore.QueuedInput> durableInputs=new ConcurrentLinkedQueue<>();
java.util.concurrent.atomic.AtomicLong inputIds=new java.util.concurrent.atomic.AtomicLong();
ConversationTurnGate gate=new ConversationTurnGate();
GoalEntity goal=new GoalEntity();
GoalSegmentRunner runner=new GoalSegmentRunner(agents,conversations,approvals,streams,new ObjectMapper(),gate,inputQueue);
@BeforeEach void setup() {
goal.setId(1L);goal.setConversationId("conv");goal.setAgentId(2L);goal.setWorkspaceId(3L);goal.setCreatedBy("alice");
ConversationEntity conv=new ConversationEntity();
conv.setConversationId("conv");conv.setAgentId(2L);conv.setWorkspaceId(3L);conv.setUsername("alice");
when(conversations.findByConversationId("conv")).thenReturn(conv);
AgentEntity agent=new AgentEntity();agent.setEnabled(true);agent.setRuntimeType("native");
when(agents.getAgent(2L)).thenReturn(agent);
when(inputQueue.enqueue(anyString(),anyLong(),anyString(),anyString(),nullable(List.class),any()))
.thenAnswer(inv -> {
var now=LocalDateTime.now();
var input=new ConversationInputQueueStore.QueuedInput(inputIds.incrementAndGet(),
inv.getArgument(0),inv.getArgument(1),inv.getArgument(2),inv.getArgument(3),
inv.getArgument(4),"queued",null,null,null,now,now);
durableInputs.add(input);
return input;
});
when(inputQueue.claimNext(anyString(),anyString(),any())).thenAnswer(inv -> {
var input=durableInputs.poll();
if(input==null) return java.util.Optional.empty();
return java.util.Optional.of(new ConversationInputQueueStore.QueuedInput(input.id(),input.conversationId(),
input.agentId(),input.createdBy(),input.message(),input.contentParts(),"claimed",
inv.getArgument(1),input.persistedMessageId(),null,input.createdAt(),LocalDateTime.now(),
input.requesterUserId(),input.selectedGoalId()));
});
when(inputQueue.bindMessage(anyLong(),anyString(),anyLong(),any())).thenReturn(true);
when(inputQueue.consume(anyLong(),anyString(),any())).thenReturn(true);
when(inputQueue.release(anyLong(),anyString(),any())).thenReturn(true);
when(inputQueue.countQueued(anyString())).thenAnswer(inv -> durableInputs.size());
MessageEntity savedUser=new MessageEntity();savedUser.setId(77L);
when(conversations.saveMessage(eq("conv"),eq("user"),anyString(),nullable(List.class),eq("queued")))
.thenReturn(savedUser);
}
private void enqueue(String message,List<MessageContentPart> parts) {
inputQueue.enqueue("conv",2L,"alice",message,parts,LocalDateTime.now());
}
@Test void persistsStreamedResultAndUsage() {
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.just(new AgentService.StreamDelta("actual output",null),
AgentService.StreamDelta.event("finish_reason",Map.of("reason","normal"))));
var result=runner.run(goal,"continue",false);
assertEquals("normal",result.finishReason());
verify(conversations).saveMessage(eq("conv"),eq("assistant"),eq("actual output"),anyList(),eq("completed"),
anyInt(),anyInt(),anyInt(),anyInt(),anyInt(),anyString(),anyString(),anyString());
assertFalse(streams.isRunning("conv"));
assertNotNull(gate.tryAcquire("conv"));
}
@Test void rejectsChangedConversationIdentityBeforeAnyModelCall() {
goal.setWorkspaceId(99L);
assertThrows(IllegalStateException.class,()->runner.run(goal,"continue",false));
verify(agents,never()).chatStructuredStream(any(),any(),any(),any(),any(),any());
}
@Test void busyConversationIsNotRegisteredOrMutated() {
var user=gate.tryAcquire("conv");
assertThrows(vip.mate.exception.MateClawException.class,()->runner.run(goal,"continue",false));
verifyNoInteractions(conversations);
user.close();
}
@Test void drainsUserInputAcceptedDuringBackgroundTurn() {
var calls=new java.util.concurrent.atomic.AtomicInteger();
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenAnswer(inv -> Flux.defer(() -> {
boolean autonomous = calls.incrementAndGet()==1;
assertTrue(GoalContinuationContext.active());
assertEquals(autonomous, GoalContinuationContext.explicitPrompt());
if(autonomous) enqueue("new user instruction",null);
return Flux.just(new AgentService.StreamDelta("output",null),
AgentService.StreamDelta.event("finish_reason",Map.of("reason","normal")));
}));
runner.run(goal,"continue",false);
assertEquals(2,calls.get());
assertEquals(0,inputQueue.countQueued("conv"));
verify(agents).chatStructuredStream(eq(2L),eq("new user instruction"),eq("conv"),eq("alice"),isNull(),any());
verify(conversations).saveMessage("conv","user","new user instruction",null,"queued");
}
@Test void managedGoalDoesNotExecuteQueuedInputSelectedForAnotherGoal() {
goal.setJsonAcceptanceRequired(true);
goal.setStatus(vip.mate.goal.model.GoalStatus.ACTIVE);
var now = LocalDateTime.now();
durableInputs.add(new ConversationInputQueueStore.QueuedInput(99L, "conv", 2L, "alice",
"instruction for another Goal", List.of(), "queued", null, null, null,
now, now, 42L, 999L));
when(agents.chatStructuredStream(eq(2L), anyString(), eq("conv"), eq("alice"), isNull(), any()))
.thenReturn(Flux.just(new AgentService.StreamDelta("incorrect execution", null)));
runner.run(goal, "continue", false);
verify(agents, never()).chatStructuredStream(any(), any(), any(), any(), any(), any());
verify(conversations).saveMessage("conv", "user", "instruction for another Goal", List.of(), "queued");
verify(conversations).saveMessage(eq("conv"),eq("assistant"),contains("was not run"),
eq(List.of()),eq("completed"));
verify(inputQueue).consume(eq(99L), anyString(), any());
}
@Test void managedGoalExecutesCurrentSelectedQueuedInput() {
goal.setJsonAcceptanceRequired(true);
goal.setStatus(vip.mate.goal.model.GoalStatus.ACTIVE);
var approvalRuns=mock(GoalApprovalRunService.class);
org.springframework.test.util.ReflectionTestUtils.setField(runner,"approvalRuns",approvalRuns);
when(approvalRuns.queuedSelectionStillCurrent(any())).thenReturn(true);
var now=LocalDateTime.now();
durableInputs.add(new ConversationInputQueueStore.QueuedInput(100L,"conv",2L,"alice",
"current Goal instruction",List.of(),"queued",null,null,null,
now,now,42L,1L));
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.just(new AgentService.StreamDelta("output",null),
AgentService.StreamDelta.event("finish_reason",Map.of("reason","normal"))));
runner.run(goal,"continue",false);
verify(agents).chatStructuredStream(eq(2L),eq("current Goal instruction"),eq("conv"),
eq("alice"),isNull(),any());
verify(inputQueue).consume(eq(100L),anyString(),any());
verify(approvalRuns).queuedSelectionStillCurrent(argThat(origin ->
origin.selectedGoalId().equals(1L) && origin.requesterUserId().equals(42L)));
}
@Test void pausedManagedGoalPreservesSelectedQueuedInputForResume() {
goal.setJsonAcceptanceRequired(true);
goal.setStatus(vip.mate.goal.model.GoalStatus.PAUSED);
var now=LocalDateTime.now();
durableInputs.add(new ConversationInputQueueStore.QueuedInput(101L,"conv",2L,"alice",
"paused Goal instruction",List.of(),"queued",null,null,null,
now,now,42L,1L));
var outcome=runner.run(goal,"continue",false);
assertInstanceOf(SegmentOutcome.Cancelled.class,outcome);
verify(agents,never()).chatStructuredStream(any(),any(),any(),any(),any(),any());
verify(conversations,never()).saveMessage("conv","user","paused Goal instruction",List.of(),"queued");
verify(inputQueue,never()).consume(eq(101L),anyString(),any());
verify(inputQueue).release(eq(101L),anyString(),any());
}
@Test void rejectedQueuedInputDoesNotBlockFollowingCurrentSelectedInput() {
goal.setJsonAcceptanceRequired(true);
goal.setStatus(vip.mate.goal.model.GoalStatus.ACTIVE);
var approvalRuns=mock(GoalApprovalRunService.class);
org.springframework.test.util.ReflectionTestUtils.setField(runner,"approvalRuns",approvalRuns);
when(approvalRuns.queuedSelectionStillCurrent(any())).thenReturn(true);
var now=LocalDateTime.now();
durableInputs.add(new ConversationInputQueueStore.QueuedInput(102L,"conv",2L,"alice",
"other Goal instruction",List.of(),"queued",null,null,null,
now,now,42L,999L));
durableInputs.add(new ConversationInputQueueStore.QueuedInput(103L,"conv",2L,"alice",
"current Goal instruction",List.of(),"queued",null,null,null,
now,now,42L,1L));
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.just(new AgentService.StreamDelta("output",null),
AgentService.StreamDelta.event("finish_reason",Map.of("reason","normal"))));
runner.run(goal,"continue",false);
verify(agents,times(1)).chatStructuredStream(eq(2L),eq("current Goal instruction"),
eq("conv"),eq("alice"),isNull(),any());
verify(inputQueue).consume(eq(102L),anyString(),any());
verify(inputQueue).consume(eq(103L),anyString(),any());
}
@Test void managedGoalDoesNotRunQueueWhenOriginalAccountIsRevoked() {
goal.setJsonAcceptanceRequired(true);
goal.setStatus(vip.mate.goal.model.GoalStatus.ACTIVE);
var approvalRuns=mock(GoalApprovalRunService.class);
org.springframework.test.util.ReflectionTestUtils.setField(runner,"approvalRuns",approvalRuns);
when(approvalRuns.queuedSelectionStillCurrent(any())).thenReturn(false);
var now=LocalDateTime.now();
durableInputs.add(new ConversationInputQueueStore.QueuedInput(104L,"conv",2L,"alice",
"revoked account instruction",List.of(),"queued",null,null,null,
now,now,42L,1L));
runner.run(goal,"continue",false);
verify(agents,never()).chatStructuredStream(any(),any(),any(),any(),any(),any());
verify(conversations).saveMessage("conv","user","revoked account instruction",List.of(),"queued");
verify(conversations).saveMessage(eq("conv"),eq("assistant"),contains("was not run"),
eq(List.of()),eq("completed"));
verify(inputQueue).consume(eq(104L),anyString(),any());
}
@Test void workerCancellationPersistsPartialEvidenceAndReleasesAdmission() throws Exception {
var subscribed=new java.util.concurrent.CountDownLatch(1);
var toolCancelled=new java.util.concurrent.atomic.AtomicBoolean();
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.concat(Flux.just(new AgentService.StreamDelta("partial evidence",null)),
Flux.<AgentService.StreamDelta>never().doOnSubscribe(s -> {
streams.registerCancellationHook("conv",()->toolCancelled.set(true));
subscribed.countDown();
})));
var failure=new java.util.concurrent.atomic.AtomicReference<Throwable>();
var result=new java.util.concurrent.atomic.AtomicReference<SegmentOutcome>();
Thread worker=Thread.ofVirtual().start(() -> {
try { result.set(runner.run(goal,"continue",false)); } catch(Throwable error) { failure.set(error); }
});
assertTrue(subscribed.await(3,java.util.concurrent.TimeUnit.SECONDS));
runner.cancel(1L);
worker.join(3000);
assertFalse(worker.isAlive());
assertTrue(toolCancelled.get(),"escaped tool process must be cancelled as well as the stream");
// Cancellation may finish the Flux before the worker receives its interrupt.
// Both paths must preserve evidence and report a stopped/interrupted outcome.
if (failure.get()==null) assertEquals("stopped",result.get().finishReason());
verify(conversations).saveMessage(eq("conv"),eq("assistant"),eq("partial evidence"),anyList(),
argThat(status -> "interrupted".equals(status) || "stopped".equals(status)),
anyInt(),anyInt(),anyInt(),anyInt(),anyInt(),anyString(),anyString(),anyString());
assertNotNull(gate.tryAcquire("conv"));
}
@Test void cancellationDoesNotInterruptDatabasePersistence() throws Exception {
var saving = new CountDownLatch(1);
var release = new CountDownLatch(1);
var interrupted = new AtomicBoolean();
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.just(new AgentService.StreamDelta("checkpoint", null)));
when(conversations.saveMessage(eq("conv"),eq("assistant"),anyString(),anyList(),anyString(),
anyInt(),anyInt(),anyInt(),anyInt(),anyInt(),anyString(),anyString(),anyString()))
.thenAnswer(inv -> {
saving.countDown();
try { release.await(3, TimeUnit.SECONDS); }
catch (InterruptedException error) { interrupted.set(true); }
return null;
});
Thread worker = Thread.ofVirtual().start(() -> runner.run(goal,"continue",false));
assertTrue(saving.await(3,TimeUnit.SECONDS));
runner.cancel(1L);
release.countDown();
worker.join(3000);
assertFalse(worker.isAlive());
assertFalse(interrupted.get(), "cancellation must not close embedded database channels by interrupting I/O");
}
@Test void shutdownPersistsQueuedInputWithAttachments() throws Exception {
var ready = new CountDownLatch(1);
MessageContentPart attachment = new MessageContentPart();
attachment.setType("file");
attachment.setPath("test-evidence.txt");
var parts = List.of(attachment);
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.<AgentService.StreamDelta>never().doOnSubscribe(s -> ready.countDown()));
Thread worker = Thread.ofVirtual().start(() -> runner.run(goal,"continue",false));
assertTrue(ready.await(3,TimeUnit.SECONDS));
enqueue("accepted steering",parts);
runner.cancelAll();
worker.join(3000);
assertFalse(worker.isAlive());
verify(conversations,never()).saveMessage("conv","user","accepted steering",parts,"queued");
assertEquals(1,inputQueue.countQueued("conv"));
}
@Test void shutdownRejectsLateWorkerAdmissionWithoutStartingModel() {
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.just(AgentService.StreamDelta.event("finish_reason",Map.of("reason","normal"))));
runner.cancelAll();
assertEquals("stopped", runner.run(goal,"continue",false).finishReason());
verify(agents,never()).chatStructuredStream(any(),any(),any(),any(),any(),any());
}
@Test void externalInterruptIsRestoredOnlyAfterCheckpointPersistence() throws Exception {
var ready = new CountDownLatch(1);
var persisted = new AtomicBoolean();
var interruptRestored = new AtomicBoolean();
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.concat(Flux.just(new AgentService.StreamDelta("partial",null)),
Flux.<AgentService.StreamDelta>never().doOnSubscribe(s -> ready.countDown())));
when(conversations.saveMessage(eq("conv"),eq("assistant"),anyString(),anyList(),anyString(),
anyInt(),anyInt(),anyInt(),anyInt(),anyInt(),anyString(),anyString(),anyString()))
.thenAnswer(inv -> {
assertFalse(Thread.currentThread().isInterrupted());
persisted.set(true);
return null;
});
Thread worker = Thread.ofVirtual().start(() -> {
try { runner.run(goal,"continue",false); }
catch (IllegalStateException expected) { interruptRestored.set(Thread.currentThread().isInterrupted()); }
});
assertTrue(ready.await(3,TimeUnit.SECONDS));
worker.interrupt();
worker.join(3000);
assertFalse(worker.isAlive());
assertTrue(persisted.get());
assertTrue(interruptRestored.get());
}
@Test void permanentFailurePersistsAcceptedQueuedInput() {
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenReturn(Flux.defer(() -> {
enqueue("user instruction",null);
return Flux.error(new IllegalArgumentException("bad config"));
}));
assertThrows(IllegalArgumentException.class,()->runner.run(goal,"continue",false));
verify(conversations,never()).saveMessage("conv","user","user instruction",null,"queued");
assertEquals(1,inputQueue.countQueued("conv"));
}
@Test void userSteeringPreservesInterruptedStatusThenRunsQueuedInput() throws Exception {
var subscribed=new java.util.concurrent.CountDownLatch(1);
var calls=new java.util.concurrent.atomic.AtomicInteger();
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenAnswer(inv -> calls.incrementAndGet()==1
? Flux.concat(Flux.just(new AgentService.StreamDelta("partial",null)),
Flux.<AgentService.StreamDelta>never().doOnSubscribe(s -> subscribed.countDown()))
: Flux.just(new AgentService.StreamDelta("answer to steering",null)));
var failure=new java.util.concurrent.atomic.AtomicReference<Throwable>();
Thread worker=Thread.ofVirtual().start(() -> {
try { runner.run(goal,"continue",false); } catch(Throwable error) { failure.set(error); }
});
assertTrue(subscribed.await(3,java.util.concurrent.TimeUnit.SECONDS));
enqueue("new instruction",null);
assertTrue(streams.requestInterrupt("conv","",2L,false));
worker.join(3000);
assertFalse(worker.isAlive());
assertNull(failure.get());
assertEquals(2,calls.get());
verify(conversations).saveMessage(eq("conv"),eq("assistant"),eq("partial"),anyList(),eq("interrupted"),
anyInt(),anyInt(),anyInt(),anyInt(),anyInt(),anyString(),anyString(),anyString());
}
@Test void goalCancellationDoesNotKillQueuedInteractiveWork() throws Exception {
var entered=new java.util.concurrent.CountDownLatch(1);
var finish=reactor.core.publisher.Sinks.<AgentService.StreamDelta>one();
var calls=new java.util.concurrent.atomic.AtomicInteger();
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenAnswer(inv -> {
if(calls.incrementAndGet()==1) {
enqueue("new question",null);
return Flux.just(AgentService.StreamDelta.event("finish_reason",Map.of("reason","normal")));
}
return finish.asMono().flux().doOnSubscribe(s -> entered.countDown());
});
var failure=new java.util.concurrent.atomic.AtomicReference<Throwable>();
Thread worker=Thread.ofVirtual().start(() -> {
try { runner.run(goal,"continue",false); } catch(Throwable error) { failure.set(error); }
});
assertTrue(entered.await(3,java.util.concurrent.TimeUnit.SECONDS));
runner.cancel(1L);
finish.tryEmitValue(AgentService.StreamDelta.event("finish_reason",Map.of("reason","normal")));
worker.join(3000);
assertFalse(worker.isAlive());
assertNull(failure.get());
}
@Test void explicitStopLatchesAcrossQueuedInputRegistrationGap() throws Exception {
var saving=new java.util.concurrent.CountDownLatch(1);
var release=new java.util.concurrent.CountDownLatch(1);
var calls=new java.util.concurrent.atomic.AtomicInteger();
when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
.thenAnswer(inv -> {
calls.incrementAndGet();
enqueue("new question",null);
return Flux.just(AgentService.StreamDelta.event("finish_reason",Map.of("reason","normal")));
});
when(conversations.saveMessage("conv","user","new question",null,"queued")).thenAnswer(inv -> {
saving.countDown();
boolean interrupted=false;
while(true) {
try { if(release.await(3,java.util.concurrent.TimeUnit.SECONDS)) break; else throw new AssertionError("release timeout"); }
catch(InterruptedException ignored) { interrupted=true; }
}
if(interrupted) Thread.currentThread().interrupt();
return null;
});
Thread worker=Thread.ofVirtual().start(() -> {
try { runner.run(goal,"continue",false); } catch(RuntimeException expected) { }
});
assertTrue(saving.await(3,java.util.concurrent.TimeUnit.SECONDS));
runner.stopConversation("conv");
release.countDown();
worker.join(3000);
assertFalse(worker.isAlive());
assertEquals(1,calls.get(),"no queued model request may start after explicit Stop");
}
}