mateclaw/mateclaw-server/src/test/java/vip/mate/goal/service/GoalSegmentRunnerTest.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.model.AgentEntity;
import vip.mate.agent.runtime.ConversationTurnGate;
import vip.mate.approval.ApprovalWorkflowService;
import vip.mate.channel.web.ChatStreamTracker;
import vip.mate.goal.model.GoalEntity;
import vip.mate.workspace.conversation.ConversationService;
import vip.mate.workspace.conversation.model.ConversationEntity;
import java.util.Map;
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());
    ConversationTurnGate gate=new ConversationTurnGate();
    GoalEntity goal=new GoalEntity();
    GoalSegmentRunner runner=new GoalSegmentRunner(agents,conversations,approvals,streams,new ObjectMapper(),gate);

    @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);
    }

    @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(() -> {
                    if(calls.incrementAndGet()==1) streams.enqueueMessage("conv","new user instruction",2L,false);
                    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());
        assertFalse(streams.hasQueuedMessage("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 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<GoalSegmentRunner.Result>();
        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 permanentFailurePersistsAcceptedQueuedInput() {
        when(agents.chatStructuredStream(eq(2L),anyString(),eq("conv"),eq("alice"),isNull(),any()))
                .thenReturn(Flux.defer(() -> {
                    streams.enqueueMessage("conv","user instruction",2L,false);
                    return Flux.error(new IllegalArgumentException("bad config"));
                }));
        assertThrows(IllegalArgumentException.class,()->runner.run(goal,"continue",false));
        verify(conversations).saveMessage("conv","user","user instruction",null,"queued");
        assertFalse(streams.hasQueuedMessage("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));
        assertTrue(streams.requestInterrupt("conv","new instruction",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) {
                        streams.enqueueMessage("conv","new question",2L,false);
                        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();
                    streams.enqueueMessage("conv","new question",2L,false);
                    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");
    }
}