test(managed-agent): Reconnect Hosted fault transport probes (#16101)

Co-authored-by: Qwen-Coder <[email protected]>

packages/managed-agent-server/verification/src/test/java/com/alibaba/qwen/code/managedagent/HostedCancellationProbe.java

@@ -33,6 +33,8 @@ final class HostedCancellationProbe implements AutoCloseable {
     private final String tenant;
     private final List<Map<String, Object>> sessions;
     private final Object service;
+    private final Object dispatcher;
+    private final Object cancellation;
     private final RuntimeTransport transport;
     private final Map<String, AtomicInteger> executions = new ConcurrentHashMap<>();
     private final Map<String, AtomicInteger> completions = new ConcurrentHashMap<>();
@@ -46,7 +48,9 @@ final class HostedCancellationProbe implements AutoCloseable {
         this.tenant = tenant;
         this.sessions = sessions;
         service = ReflectionTestUtils.getField(broker, "service");
-        transport = (RuntimeTransport) ReflectionTestUtils.getField(service, "transport");
+        dispatcher = ReflectionTestUtils.getField(service, "executionDispatcher");
+        cancellation = ReflectionTestUtils.getField(service, "executionCancellation");
+        transport = (RuntimeTransport) ReflectionTestUtils.getField(dispatcher, "transport");
         RuntimeTransport observed = (RuntimeTransport) Proxy.newProxyInstance(RuntimeTransport.class.getClassLoader(),
                 new Class<?>[] {RuntimeTransport.class}, (proxy, method, args) -> {
                     String session = args != null && args.length > 1 && args[1] instanceof RuntimeSession runtime
@@ -71,7 +75,8 @@ final class HostedCancellationProbe implements AutoCloseable {
                     }
                     return result;
                 });
-        ReflectionTestUtils.setField(service, "transport", observed);
+        ReflectionTestUtils.setField(dispatcher, "transport", observed);
+        ReflectionTestUtils.setField(cancellation, "transport", observed);
         server.createContext("/cancellation/", exchange -> {
             try {
                 assertThat(exchange.getRequestMethod()).isEqualTo("POST");
@@ -258,7 +263,8 @@ final class HostedCancellationProbe implements AutoCloseable {
 
     @Override
     public void close() throws InterruptedException {
-        ReflectionTestUtils.setField(service, "transport", transport);
+        ReflectionTestUtils.setField(dispatcher, "transport", transport);
+        ReflectionTestUtils.setField(cancellation, "transport", transport);
         Object provisioner = ReflectionTestUtils.getField(service, "provisioner");
         Object local = ReflectionTestUtils.getField(provisioner, "delegate");
         var owned = (Map<?, ?>) ReflectionTestUtils.getField(local, "owned");

packages/managed-agent-server/verification/src/test/java/com/alibaba/qwen/code/managedagent/HostedExecutionTransportProbeTest.java

@@ -0,0 +1,143 @@
+package com.alibaba.qwen.code.managedagent;
+
+import com.alibaba.qwen.code.managedagent.config.runtime.EmbeddedRuntimeBroker;
+import com.alibaba.qwen.code.runtimebroker.RuntimeLease;
+import com.alibaba.qwen.code.runtimebroker.RuntimeProvisionRequest;
+import com.alibaba.qwen.code.runtimebroker.RuntimeProvisioner;
+import com.alibaba.qwen.code.runtimebroker.RuntimeScope;
+import com.alibaba.qwen.code.runtimebroker.RuntimeTransport;
+import com.alibaba.qwen.code.runtimebroker.ToolExecutionRecord;
+import com.alibaba.qwen.code.runtimebroker.persistence.InMemoryRuntimeBindingRepository;
+import com.alibaba.qwen.code.runtimebroker.persistence.InMemoryRuntimeSessionRepository;
+import com.alibaba.qwen.code.runtimebroker.persistence.InMemoryToolExecutionRepository;
+import com.alibaba.qwen.code.runtimebroker.service.RuntimeBrokerService;
+import com.sun.net.httpserver.HttpServer;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import java.net.InetSocketAddress;
+import java.net.URI;
+import java.time.Duration;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+@Timeout(10)
+class HostedExecutionTransportProbeTest {
+    @Test
+    void cancellationProbeObservesExecutionAndCancellationAndRestoresBothOwners() throws Exception {
+        try (var fixture = new Fixture()) {
+            HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+            var probe = new HostedCancellationProbe(mock(JdbcTemplate.class), "tenant", List.of(),
+                    fixture.embedded, server);
+            try {
+                ToolExecutionRecord execution = fixture.execute("first");
+                assertThat(execution.getState()).isEqualTo(ToolExecutionRecord.State.EXECUTING);
+                assertThat(count(probe, "executions")).isEqualTo(1);
+                assertThat(count(probe, "completions")).isZero();
+
+                assertThat(fixture.service.cancelExecution("harness", "runtime", execution.getExecutionCallId())
+                        .toCompletableFuture().join().getState()).isEqualTo(ToolExecutionRecord.State.CANCEL_REQUESTED);
+                assertThat(count(probe, "cancellations")).isEqualTo(1);
+                assertThat(count(probe, "completions")).isZero();
+                fixture.answer.complete(Map.of("executionStatus", "cancelled"));
+                assertThat(fixture.state(execution)).isEqualTo(ToolExecutionRecord.State.SETTLED);
+                assertThat(count(probe, "completions")).isEqualTo(1);
+            } finally {
+                probe.close();
+                server.stop(0);
+            }
+
+            fixture.answer = new CompletableFuture<>();
+            ToolExecutionRecord afterClose = fixture.execute("after-close");
+            fixture.service.cancelExecution("harness", "runtime", afterClose.getExecutionCallId())
+                    .toCompletableFuture().join();
+            fixture.answer.complete(Map.of("executionStatus", "cancelled"));
+            assertThat(fixture.state(afterClose)).isEqualTo(ToolExecutionRecord.State.SETTLED);
+            assertThat(count(probe, "executions")).isEqualTo(1);
+            assertThat(count(probe, "cancellations")).isEqualTo(1);
+            assertThat(count(probe, "completions")).isEqualTo(1);
+        }
+    }
+
+    @Test
+    void statusReplyLossGateHoldsSettlementUntilRelease() throws Exception {
+        try (var fixture = new Fixture()) {
+            var gate = new CompletableFuture<Void>();
+            HostedWorkspaceToolTurnIT.holdExecutionResults(fixture.embedded, "harness", gate);
+            ToolExecutionRecord execution = fixture.execute("held");
+            fixture.answer.complete(Map.of("executionStatus", "success"));
+            assertThat(gate.isDone()).isFalse();
+            assertThat(fixture.state(execution)).isEqualTo(ToolExecutionRecord.State.EXECUTING);
+            gate.complete(null);
+            assertThat(fixture.state(execution)).isEqualTo(ToolExecutionRecord.State.SETTLED);
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private static int count(HostedCancellationProbe probe, String field) {
+        var counts = (Map<String, AtomicInteger>) ReflectionTestUtils.getField(probe, field);
+        AtomicInteger value = counts.get("harness");
+        return value == null ? 0 : value.get();
+    }
+
+    private static final class Fixture implements AutoCloseable {
+        private final RuntimeTransport transport = mock(RuntimeTransport.class);
+        private final EmbeddedRuntimeBroker embedded = mock(EmbeddedRuntimeBroker.class);
+        private final RuntimeBrokerService service;
+        private CompletableFuture<Map<String, Object>> answer = new CompletableFuture<>();
+
+        private Fixture() {
+            RuntimeScope scope = new RuntimeScope("tenant", "workspace", "1", "/workspace", "capability", "workspace");
+            when(transport.acquire(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+            when(transport.reclaimMount(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+            when(transport.execute(any(), any(), any())).thenAnswer(ignored -> answer);
+            when(transport.cancel(any(), any(), any()))
+                    .thenReturn(CompletableFuture.completedFuture(Map.of("state", "cancel_requested")));
+            service = new RuntimeBrokerService(ignored -> CompletableFuture.completedFuture(scope),
+                    new Provisioner(), transport, new InMemoryRuntimeBindingRepository(),
+                    new InMemoryRuntimeSessionRepository(), new InMemoryToolExecutionRepository(),
+                    "broker", Duration.ofMinutes(1), Duration.ofMinutes(1));
+            ReflectionTestUtils.setField(embedded, "service", service);
+            service.acquire("harness", "runtime", "bootstrap").toCompletableFuture().join();
+        }
+
+        private ToolExecutionRecord execute(String key) {
+            return service.createExecution("harness", "runtime", key,
+                    Map.of("sessionId", "runtime", "promptId", "prompt", "callId", key, "argsDigest", "digest"))
+                    .toCompletableFuture().join();
+        }
+
+        private ToolExecutionRecord.State state(ToolExecutionRecord execution) {
+            return service.getExecution("harness", "runtime", execution.getExecutionCallId())
+                    .toCompletableFuture().join().getState();
+        }
+
+        @Override
+        public void close() {
+            service.close();
+        }
+    }
+
+    private static final class Provisioner implements RuntimeProvisioner {
+        private final OwnedProcesses delegate = new OwnedProcesses(Map.of());
+
+        @Override
+        public CompletionStage<RuntimeLease> provision(RuntimeProvisionRequest request) {
+            return CompletableFuture.completedFuture(new RuntimeLease("instance", URI.create("http://127.0.0.1:1234"),
+                    "token", "lease", 1));
+        }
+    }
+
+    private record OwnedProcesses(Map<String, Object> owned) {
+    }
+}

packages/managed-agent-server/verification/src/test/java/com/alibaba/qwen/code/managedagent/HostedWorkspaceToolTurnIT.java

@@ -254,6 +254,23 @@ class HostedWorkspaceToolTurnIT {
         });
     }
 
+    static void holdExecutionResults(EmbeddedRuntimeBroker broker, String harnessSessionId,
+            CompletableFuture<Void> gate) {
+        Object service = ReflectionTestUtils.getField(broker, "service");
+        Object dispatcher = ReflectionTestUtils.getField(service, "executionDispatcher");
+        RuntimeTransport transport = (RuntimeTransport) ReflectionTestUtils.getField(dispatcher, "transport");
+        RuntimeTransport gated = (RuntimeTransport) Proxy.newProxyInstance(RuntimeTransport.class.getClassLoader(),
+                new Class<?>[] {RuntimeTransport.class}, (proxy, method, args) -> {
+                    Object result = method.invoke(transport, args);
+                    if (method.getName().equals("execute")
+                            && ((RuntimeSession) args[1]).getHarnessSessionId().equals(harnessSessionId)) {
+                        return ((CompletionStage<?>) result).thenCombine(gate, (value, ignored) -> value);
+                    }
+                    return result;
+                });
+        ReflectionTestUtils.setField(dispatcher, "transport", gated);
+    }
+
     private void runDriver(List<String> cases, String driverName) throws Exception {
         boolean latency = driverName.equals("latency");
         boolean unicode = driverName.equals("shell-unicode");
@@ -424,18 +441,7 @@ class HostedWorkspaceToolTurnIT {
             });
             if (faults && cases.contains("status")) {
                 String statusSession = sessions.get(cases.indexOf("status")).get("sessionId").toString();
-                Object service = ReflectionTestUtils.getField(broker, "service");
-                RuntimeTransport transport = (RuntimeTransport) ReflectionTestUtils.getField(service, "transport");
-                RuntimeTransport gated = (RuntimeTransport) Proxy.newProxyInstance(RuntimeTransport.class.getClassLoader(),
-                        new Class<?>[] {RuntimeTransport.class}, (proxy, method, args) -> {
-                            Object result = method.invoke(transport, args);
-                            if (method.getName().equals("execute")
-                                    && ((RuntimeSession) args[1]).getHarnessSessionId().equals(statusSession)) {
-                                return ((CompletionStage<?>) result).thenCombine(statusGate, (value, ignored) -> value);
-                            }
-                            return result;
-                        });
-                ReflectionTestUtils.setField(service, "transport", gated);
+                holdExecutionResults(broker, statusSession, statusGate);
             }
             HostedCancellationProbe cancellationProbe = cancellations
                     ? new HostedCancellationProbe(jdbc, tenant, sessions, broker, gateServer) : null;