修复SuperAgent自动分发恢复逻辑

2026-07-13 00:35:43 +08:00
parent 53b530d0b3
commit 50180f88cb
7 changed files with 289 additions and 16 deletions
--- a/server/src/main/java/cn/nianxx/thhotel/integrations/ai/superagent/repository/MybatisSuperAgentDispatchRunRepository.java
+++ b/server/src/main/java/cn/nianxx/thhotel/integrations/ai/superagent/repository/MybatisSuperAgentDispatchRunRepository.java
@@ -47,7 +47,7 @@ public class MybatisSuperAgentDispatchRunRepository implements SuperAgentDispatc
    }

    /**
-     * 抢占 PENDING/RETRYABLE_FAILED 记录。先查候选再逐条条件更新，降低多 worker 重复执行概率。
+     * 抢占 PENDING/RETRYABLE_FAILED 或锁已过期的 RUNNING 记录。先查候选再逐条条件更新，降低多 worker 重复执行概率。
     */
    @Override
    public List<SuperAgentDispatchRunSnapshot> claimDue(
--- a/server/src/main/java/cn/nianxx/thhotel/integrations/ai/superagent/service/impl/SuperAgentDispatchServiceImpl.java
+++ b/server/src/main/java/cn/nianxx/thhotel/integrations/ai/superagent/service/impl/SuperAgentDispatchServiceImpl.java
@@ -80,7 +80,6 @@ public class SuperAgentDispatchServiceImpl implements SuperAgentDispatchService
            return;
        }
        if (captureResult.inboxId() == null
-                || !captureResult.created()
                || !CAPTURE_STATUS_RECEIVED.equals(captureResult.captureStatus())) {
            return;
        }
@@ -116,7 +115,10 @@ public class SuperAgentDispatchServiceImpl implements SuperAgentDispatchService
        LocalDateTime now = nowUtc();
        List<SuperAgentDispatchRunSnapshot> runs = runRepository.claimDue(
                "local-worker",
-                List.of(SuperAgentDispatchStatus.PENDING.code(), SuperAgentDispatchStatus.RETRYABLE_FAILED.code()),
+                List.of(
+                        SuperAgentDispatchStatus.PENDING.code(),
+                        SuperAgentDispatchStatus.RETRYABLE_FAILED.code(),
+                        SuperAgentDispatchStatus.RUNNING.code()),
                now,
                now.plus(lockTtl()),
                safeBatchSize());
--- a/server/src/main/java/cn/nianxx/thhotel/integrations/ai/superagent/service/impl/SuperAgentOpenApiClientImpl.java
+++ b/server/src/main/java/cn/nianxx/thhotel/integrations/ai/superagent/service/impl/SuperAgentOpenApiClientImpl.java
@@ -7,17 +7,18 @@ import cn.nianxx.thhotel.integrations.ai.superagent.common.result.SuperAgentOpen
 import cn.nianxx.thhotel.integrations.ai.superagent.service.SuperAgentOpenApiClient;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import java.io.InputStream;
 import java.io.IOException;
-import java.nio.charset.StandardCharsets;
+import java.io.InputStream;
+import java.io.InputStreamReader;
 import java.net.URI;
+import java.nio.charset.StandardCharsets;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
-import java.io.InputStreamReader;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.function.Consumer;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;

 /**
@@ -29,17 +30,31 @@ public class SuperAgentOpenApiClientImpl implements SuperAgentOpenApiClient {
    private final SuperAgentOpenApiProperties properties;
    private final SuperAgentOpenApiSseParser sseParser;
    private final ObjectMapper objectMapper;
+    private final HttpClient injectedHttpClient;

    /**
     * 注入 SuperAgent 配置、SSE 解析器和 JSON 工具。
     */
+    @Autowired
    public SuperAgentOpenApiClientImpl(
            SuperAgentOpenApiProperties properties,
            SuperAgentOpenApiSseParser sseParser,
            ObjectMapper objectMapper) {
+        this(properties, sseParser, objectMapper, null);
+    }
+
+    /**
+     * 注入可替换的 HTTP 客户端，供单元测试验证低层 HTTP 边界行为。
+     */
+    SuperAgentOpenApiClientImpl(
+            SuperAgentOpenApiProperties properties,
+            SuperAgentOpenApiSseParser sseParser,
+            ObjectMapper objectMapper,
+            HttpClient injectedHttpClient) {
        this.properties = properties;
        this.sseParser = sseParser;
        this.objectMapper = objectMapper;
+        this.injectedHttpClient = injectedHttpClient;
    }

    /**
@@ -129,7 +144,7 @@ public class SuperAgentOpenApiClientImpl implements SuperAgentOpenApiClient {
        HttpResponse<InputStream> httpResponse = httpClient.send(
                httpRequest,
                HttpResponse.BodyHandlers.ofInputStream());
-        require2xx(httpResponse.statusCode(), null);
+        require2xxOrClose(httpResponse.statusCode(), httpResponse.body());
        SuperAgentOpenApiSseParser.ParsedState state = sseParser.newState(sessionId);
        applyRunLocation(baseUri, httpResponse, state);
        try {
@@ -160,6 +175,9 @@ public class SuperAgentOpenApiClientImpl implements SuperAgentOpenApiClient {
     * 构造 JDK HttpClient；SSE 请求不设置整体 read timeout，避免长任务被固定超时截断。
     */
    private HttpClient httpClient() {
+        if (injectedHttpClient != null) {
+            return injectedHttpClient;
+        }
        return HttpClient.newBuilder()
                .connectTimeout(properties.getConnectTimeout())
                .version(HttpClient.Version.HTTP_1_1)
@@ -202,6 +220,31 @@ public class SuperAgentOpenApiClientImpl implements SuperAgentOpenApiClient {
        }
    }

+    /**
+     * 校验流式响应状态码；失败时必须关闭响应流，避免连接资源泄漏。
+     */
+    private void require2xxOrClose(int statusCode, InputStream responseBody) {
+        if (statusCode >= 200 && statusCode < 300) {
+            return;
+        }
+        closeQuietly(responseBody);
+        require2xx(statusCode, null);
+    }
+
+    /**
+     * 静默关闭错误响应流；关闭异常不覆盖原始 HTTP 状态错误。
+     */
+    private void closeQuietly(InputStream responseBody) {
+        if (responseBody == null) {
+            return;
+        }
+        try {
+            responseBody.close();
+        } catch (IOException exception) {
+            // 错误响应流关闭失败不覆盖原始 HTTP 状态错误。
+        }
+    }
+
    /**
     * 从 Content-Location 保存 run URI 和 runId，供 EOF 后恢复使用。
     */
@@ -297,7 +340,7 @@ public class SuperAgentOpenApiClientImpl implements SuperAgentOpenApiClient {
        HttpResponse<InputStream> response = httpClient.send(
                builder.build(),
                HttpResponse.BodyHandlers.ofInputStream());
-        require2xx(response.statusCode(), null);
+        require2xxOrClose(response.statusCode(), response.body());
        consumeResponse(response.body(), state, traceConsumer);
    }

--- a/server/src/test/java/cn/nianxx/thhotel/integrations/ai/superagent/service/impl/SuperAgentDispatchServiceImplTest.java
+++ b/server/src/test/java/cn/nianxx/thhotel/integrations/ai/superagent/service/impl/SuperAgentDispatchServiceImplTest.java
@@ -2,6 +2,7 @@ package cn.nianxx.thhotel.integrations.ai.superagent.service.impl;

 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
@@ -66,13 +67,46 @@ class SuperAgentDispatchServiceImplTest {
                true,
                false,
                "FAILED"));
+
+        verify(runRepository, never()).insertIfAbsent(any(SuperAgentDispatchRunDraft.class));
+    }
+
+    @Test
+    void shouldCreateDispatchForDuplicateReceivedSourceMessageToCompensateMissingOutbox() {
+        SuperAgentDispatchRunRepository runRepository = mock(SuperAgentDispatchRunRepository.class);
+        when(runRepository.insertIfAbsent(any(SuperAgentDispatchRunDraft.class))).thenReturn(99002L);
+        SuperAgentDispatchServiceImpl service = service(runRepository, mock(SourceMessageInboxRepository.class),
+                mock(SuperAgentOpenApiClient.class), properties(true, false));
+
        service.enqueueAgentBusRealtime(agentBusCommand("AGENTBUS"), new SourceMessageCaptureResult(
                88004L,
                false,
                false,
                "RECEIVED"));

-        verify(runRepository, never()).insertIfAbsent(any(SuperAgentDispatchRunDraft.class));
+        ArgumentCaptor<SuperAgentDispatchRunDraft> captor = ArgumentCaptor.forClass(SuperAgentDispatchRunDraft.class);
+        verify(runRepository).insertIfAbsent(captor.capture());
+        assertThat(captor.getValue().sourceMessageId()).isEqualTo(88004L);
+        assertThat(captor.getValue().dispatchSource()).isEqualTo(SuperAgentDispatchSource.AGENTBUS_REALTIME.code());
+    }
+
+    @Test
+    void shouldClaimExpiredRunningDispatchesForCrashRecovery() {
+        SuperAgentDispatchRunRepository runRepository = mock(SuperAgentDispatchRunRepository.class);
+        when(runRepository.claimDue(any(), any(), any(), any(), anyInt())).thenReturn(List.of());
+        SuperAgentDispatchServiceImpl service = service(runRepository, mock(SourceMessageInboxRepository.class),
+                mock(SuperAgentOpenApiClient.class), properties(true, true));
+
+        service.processDueDispatches();
+
+        @SuppressWarnings("unchecked")
+        ArgumentCaptor<List<String>> statusesCaptor = ArgumentCaptor.forClass(List.class);
+        verify(runRepository).claimDue(any(), statusesCaptor.capture(), any(), any(), anyInt());
+        assertThat(statusesCaptor.getValue())
+                .contains(
+                        SuperAgentDispatchStatus.PENDING.code(),
+                        SuperAgentDispatchStatus.RETRYABLE_FAILED.code(),
+                        SuperAgentDispatchStatus.RUNNING.code());
    }

    @Test
--- a/server/src/test/java/cn/nianxx/thhotel/integrations/ai/superagent/service/impl/SuperAgentOpenApiClientImplTest.java
+++ b/server/src/test/java/cn/nianxx/thhotel/integrations/ai/superagent/service/impl/SuperAgentOpenApiClientImplTest.java
@@ -1,22 +1,39 @@
 package cn.nianxx.thhotel.integrations.ai.superagent.service.impl;

 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;

 import cn.nianxx.thhotel.integrations.ai.superagent.common.request.SuperAgentMailDebugRequest;
 import cn.nianxx.thhotel.integrations.ai.superagent.common.result.SuperAgentOpenApiResult;
 import cn.nianxx.thhotel.integrations.ai.superagent.common.result.SuperAgentOpenApiTraceEvent;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.sun.net.httpserver.HttpServer;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
 import java.io.OutputStream;
+import java.net.Authenticator;
+import java.net.CookieHandler;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
+import java.net.ProxySelector;
+import java.net.http.HttpClient;
+import java.net.http.HttpHeaders;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
 import java.nio.charset.StandardCharsets;
+import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.SSLParameters;
+import javax.net.ssl.SSLSession;
 import org.junit.jupiter.api.Test;

 class SuperAgentOpenApiClientImplTest {
@@ -180,4 +197,178 @@ class SuperAgentOpenApiClientImplTest {
            server.stop(0);
        }
    }
+
+    @Test
+    void shouldCloseStreamingResponseBodyWhenStatusIsNot2xx() {
+        TrackingInputStream responseBody = new TrackingInputStream("bad gateway".getBytes(StandardCharsets.UTF_8));
+        FakeHttpClient httpClient = new FakeHttpClient(responseBody);
+        SuperAgentOpenApiClientImpl client = new SuperAgentOpenApiClientImpl(
+                openApiProperties(),
+                new SuperAgentOpenApiSseParser(new ObjectMapper()),
+                new ObjectMapper(),
+                httpClient);
+
+        assertThatThrownBy(() -> client.invokeMailDebug(new SuperAgentMailDebugRequest(
+                "debug message",
+                "debug-500-idempotency",
+                Map.of("source", "unit-test"))))
+                .isInstanceOf(SuperAgentOpenApiException.class)
+                .hasMessageContaining("status=502");
+        assertThat(responseBody.closed()).isTrue();
+        assertThat(httpClient.streamPostCount()).isEqualTo(1);
+    }
+
+    private SuperAgentOpenApiProperties openApiProperties() {
+        SuperAgentOpenApiProperties properties = new SuperAgentOpenApiProperties();
+        properties.setEnabled(true);
+        properties.setBaseUrl("http://superagent.test");
+        properties.setApiKey("df_open_test");
+        properties.setExternalSubjectId("debug-subject");
+        return properties;
+    }
+
+    private static final class TrackingInputStream extends ByteArrayInputStream {
+        private boolean closed;
+
+        private TrackingInputStream(byte[] buffer) {
+            super(buffer);
+        }
+
+        @Override
+        public void close() throws IOException {
+            closed = true;
+            super.close();
+        }
+
+        private boolean closed() {
+            return closed;
+        }
+    }
+
+    private static final class FakeHttpClient extends HttpClient {
+        private final TrackingInputStream streamingBody;
+        private final AtomicInteger streamPostCount = new AtomicInteger();
+
+        private FakeHttpClient(TrackingInputStream streamingBody) {
+            this.streamingBody = streamingBody;
+        }
+
+        private int streamPostCount() {
+            return streamPostCount.get();
+        }
+
+        @Override
+        public Optional<CookieHandler> cookieHandler() {
+            return Optional.empty();
+        }
+
+        @Override
+        public Optional<Duration> connectTimeout() {
+            return Optional.empty();
+        }
+
+        @Override
+        public Redirect followRedirects() {
+            return Redirect.NEVER;
+        }
+
+        @Override
+        public Optional<ProxySelector> proxy() {
+            return Optional.empty();
+        }
+
+        @Override
+        public SSLContext sslContext() {
+            try {
+                return SSLContext.getDefault();
+            } catch (NoSuchAlgorithmException exception) {
+                throw new IllegalStateException(exception);
+            }
+        }
+
+        @Override
+        public SSLParameters sslParameters() {
+            return new SSLParameters();
+        }
+
+        @Override
+        public Optional<Authenticator> authenticator() {
+            return Optional.empty();
+        }
+
+        @Override
+        public Version version() {
+            return Version.HTTP_1_1;
+        }
+
+        @Override
+        public Optional<Executor> executor() {
+            return Optional.empty();
+        }
+
+        @Override
+        public <T> HttpResponse<T> send(HttpRequest request, HttpResponse.BodyHandler<T> responseBodyHandler) {
+            String path = request.uri().getPath();
+            if (path.endsWith("/api/open/agent-sessions")) {
+                return response(request, 200, "{\"session_id\":\"session-http-500\"}");
+            }
+            if (path.endsWith("/api/open/agent-sessions/session-http-500/messages/stream")) {
+                streamPostCount.incrementAndGet();
+                return response(request, 502, streamingBody);
+            }
+            throw new AssertionError("Unexpected request path: " + path);
+        }
+
+        @Override
+        public <T> CompletableFuture<HttpResponse<T>> sendAsync(
+                HttpRequest request,
+                HttpResponse.BodyHandler<T> responseBodyHandler) {
+            return CompletableFuture.failedFuture(new UnsupportedOperationException());
+        }
+
+        @Override
+        public <T> CompletableFuture<HttpResponse<T>> sendAsync(
+                HttpRequest request,
+                HttpResponse.BodyHandler<T> responseBodyHandler,
+                HttpResponse.PushPromiseHandler<T> pushPromiseHandler) {
+            return CompletableFuture.failedFuture(new UnsupportedOperationException());
+        }
+
+        @SuppressWarnings("unchecked")
+        private <T> HttpResponse<T> response(HttpRequest request, int statusCode, Object body) {
+            return new SimpleHttpResponse<>(request, statusCode, (T) body);
+        }
+    }
+
+    private record SimpleHttpResponse<T>(
+            HttpRequest request,
+            int statusCode,
+            T body
+    ) implements HttpResponse<T> {
+
+        @Override
+        public Optional<HttpResponse<T>> previousResponse() {
+            return Optional.empty();
+        }
+
+        @Override
+        public HttpHeaders headers() {
+            return HttpHeaders.of(Map.of(), (name, value) -> true);
+        }
+
+        @Override
+        public Optional<SSLSession> sslSession() {
+            return Optional.empty();
+        }
+
+        @Override
+        public java.net.URI uri() {
+            return request.uri();
+        }
+
+        @Override
+        public HttpClient.Version version() {
+            return HttpClient.Version.HTTP_1_1;
+        }
+    }
 }