This is an automated email from the ASF dual-hosted git repository.

Aias00 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shenyu.git


The following commit(s) were added to refs/heads/master by this push:
     new 9aa6bee969 OSPP2:Preserve MCP timeout responses and metrics (#7574)
9aa6bee969 is described below

commit 9aa6bee9697d3ea378d9bcbf86e81c4a56e1d870
Author: BobSong <[email protected]>
AuthorDate: Sat Oct 17 20:48:06 2026 +0800

    OSPP2:Preserve MCP timeout responses and metrics (#7574)
    
    * fix(agent-gateway): preserve MCP timeout responses and metrics
    
    * docs(agent-gateway): explain HTTP 504 timeout classification
    
    ---------
    
    Co-authored-by: BobSong-dev <[email protected]>
---
 .../gateway/protocol/AgentMcpCallObservation.java  |   4 +-
 .../agent/gateway/protocol/AgentMcpDispatcher.java |  10 ++
 .../gateway/remote/RequestScopedMcpClient.java     |   3 +-
 .../gateway/protocol/AgentMcpAggregationTest.java  |  23 ++++
 .../protocol/AgentMcpCallObservationTest.java      |  30 ++++++
 .../gateway/remote/RequestScopedMcpClientTest.java | 120 +++++++++++++++++++++
 6 files changed, 188 insertions(+), 2 deletions(-)

diff --git 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpCallObservation.java
 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpCallObservation.java
index 4247687b44..ae9c8bf837 100644
--- 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpCallObservation.java
+++ 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpCallObservation.java
@@ -101,8 +101,10 @@ final class AgentMcpCallObservation {
         outcome = response.path("result").path("isError").asBoolean(false) ? 
Outcome.TOOL_ERROR : Outcome.SUCCESS;
     }
 
+    /** HTTP 504 is intentionally folded into {@link Outcome#TIMEOUT}, not a 
server error. */
     void failure(final Throwable error) {
-        if (error instanceof TimeoutException) {
+        if (error instanceof TimeoutException
+                || error instanceof AgentMcpProtocolException && 
((AgentMcpProtocolException) error).getHttpStatus() == 504) {
             outcome = Outcome.TIMEOUT;
         } else if (error instanceof AgentMcpProtocolException && 
((AgentMcpProtocolException) error).getHttpStatus() < 500) {
             outcome = Outcome.REJECTED;
diff --git 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpDispatcher.java
 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpDispatcher.java
index 20c4b63a68..759a36664b 100644
--- 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpDispatcher.java
+++ 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpDispatcher.java
@@ -33,6 +33,8 @@ import 
org.apache.shenyu.plugin.agent.gateway.tool.AgentToolRegistry;
 import reactor.core.publisher.Mono;
 
 import java.util.Map;
+import java.net.http.HttpTimeoutException;
+import java.util.concurrent.TimeoutException;
 import java.util.Objects;
 import java.util.function.Supplier;
 
@@ -87,6 +89,14 @@ public final class AgentMcpDispatcher {
             }
             return route(request, context).contextWrite(reactorContext -> 
reactorContext.put(AgentMcpExecutionContext.class, context));
         }).onErrorMap(error -> {
+            if (error instanceof TimeoutException) {
+                return error;
+            }
+            if (error instanceof HttpTimeoutException) {
+                TimeoutException timeout = new TimeoutException("Upstream 
request timed out");
+                timeout.initCause(error);
+                return timeout;
+            }
             if (error instanceof AgentMcpProtocolException) {
                 return error;
             }
diff --git 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/remote/RequestScopedMcpClient.java
 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/remote/RequestScopedMcpClient.java
index 33c1699d1d..1d5a75001d 100644
--- 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/remote/RequestScopedMcpClient.java
+++ 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/main/java/org/apache/shenyu/plugin/agent/gateway/remote/RequestScopedMcpClient.java
@@ -46,6 +46,7 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Flow;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -192,7 +193,7 @@ public final class RequestScopedMcpClient implements 
RemoteMcpEndpoint {
             Instant deadline = transportContext.get("deadline") instanceof 
Instant value ? value : Instant.now().plusSeconds(4);
             Duration budget = Duration.between(Instant.now(), deadline);
             if (budget.isNegative() || budget.isZero()) {
-                return Mono.error(new IllegalStateException("Deadline expired 
before upstream side effect"));
+                return Mono.error(new TimeoutException("Deadline expired 
before upstream side effect"));
             }
             byte[] requestBytes = 
message.toString().getBytes(StandardCharsets.UTF_8);
             if (requestBytes.length > MAX_REQUEST_BYTES) {
diff --git 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpAggregationTest.java
 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpAggregationTest.java
index 323be23673..ae038301dd 100644
--- 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpAggregationTest.java
+++ 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpAggregationTest.java
@@ -34,11 +34,14 @@ import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 
 import java.nio.charset.StandardCharsets;
+import java.net.http.HttpTimeoutException;
+import java.net.http.HttpConnectTimeoutException;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Function;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -54,6 +57,23 @@ class AgentMcpAggregationTest {
 
     private final AtomicInteger localCalls = new AtomicInteger();
 
+    private Throwable remoteFailure;
+
+    @ParameterizedTest
+    @ValueSource(ints = {0, 1, 2})
+    void remoteTimeoutIsNotAnInternalErrorAndReleasesItsLease(final int type) {
+        remoteFailure = switch (type) {
+            case 0 -> new TimeoutException("private-reactor-timeout");
+            case 1 -> new HttpTimeoutException("private-upstream-timeout");
+            default -> new 
HttpConnectTimeoutException("private-connect-timeout");
+        };
+        StepVerifier.create(dispatcher("local").dispatch(request("tools/call", 
"orders.lookup"), () -> context("A", Set.of("orders.lookup"))))
+                .expectErrorMatches(error -> error instanceof 
TimeoutException).verify();
+        assertEquals(1, calls.get());
+        assertEquals(0, active.get());
+        assertEquals(0, localCalls.get());
+    }
+
     @ParameterizedTest
     @ValueSource(strings = {"server/discover", "tools/list", "tools/call"})
     void 
shouldDistinguishUnavailableCatalogWithoutFallingBackToLocalTools(final String 
method) {
@@ -207,6 +227,9 @@ class AgentMcpAggregationTest {
             @Override
             public Mono<ObjectNode> invoke(final String name, final JsonObject 
arguments, final AgentMcpExecutionContext context) {
                 calls.incrementAndGet();
+                if (Objects.nonNull(remoteFailure)) {
+                    return Mono.error(remoteFailure);
+                }
                 ObjectNode result = mapper.createObjectNode().put("isError", 
true).put("extension", true);
                 result.putObject("structuredContent").put("subject", 
context.getSubject());
                 result.putArray("content").addObject().put("type", 
"text").put("text", "Business failure");
diff --git 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpCallObservationTest.java
 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpCallObservationTest.java
index ccad4dad40..a154ee8125 100644
--- 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpCallObservationTest.java
+++ 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/protocol/AgentMcpCallObservationTest.java
@@ -46,6 +46,8 @@ import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 
 import java.nio.charset.StandardCharsets;
+import java.net.http.HttpTimeoutException;
+import java.net.http.HttpConnectTimeoutException;
 import java.time.Duration;
 import java.util.List;
 import java.util.Objects;
@@ -69,6 +71,34 @@ class AgentMcpCallObservationTest {
 
     private final AgentTrafficContext traffic = new 
AgentTrafficContext("private-id", "mcp", "selector", "rule");
 
+    @ParameterizedTest
+    @CsvSource({"json,0", "sse,0", "json,1", "sse,1", "json,2", "sse,2"})
+    void innerTimeoutIsAlways504AndRecordedOnce(final String mode, final int 
type) {
+        Throwable failure = switch (type) {
+            case 0 -> new TimeoutException("private-timeout");
+            case 1 -> new HttpTimeoutException("private-upstream-timeout");
+            default -> new 
HttpConnectTimeoutException("private-connect-timeout");
+        };
+        Recording observer = new Recording();
+        MockServerWebExchange exchange = exchange("tools/call", observer);
+        StepVerifier.create(handler(tool(Mono.error(failure)), 
Set.of("read")).handle(exchange, config(mode, 30000), traffic, 
1)).verifyComplete();
+        assertEquals(504, exchange.getResponse().getStatusCode().value());
+        assertEquals(MediaType.APPLICATION_JSON, 
exchange.getResponse().getHeaders().getContentType());
+        assertEquals("{\"error\":\"Request deadline exceeded\"}", 
exchange.getResponse().getBodyAsString().block());
+        observer.assertOnly(Outcome.TIMEOUT);
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"json", "sse"})
+    void explicitGatewayTimeoutIsNotRecordedAsServerError(final String mode) {
+        Recording observer = new Recording();
+        MockServerWebExchange exchange = exchange("tools/call", observer);
+        AgentMcpProtocolException failure = new AgentMcpProtocolException(504, 
-32001, "Remote deadline exceeded", null, null);
+        StepVerifier.create(handler(tool(Mono.error(failure)), 
Set.of("read")).handle(exchange, config(mode, 30000), traffic, 
1)).verifyComplete();
+        assertEquals(504, exchange.getResponse().getStatusCode().value());
+        observer.assertOnly(Outcome.TIMEOUT);
+    }
+
     @ParameterizedTest
     @CsvSource({"json,SUCCESS", "sse,SUCCESS", "json,TOOL_ERROR", 
"sse,TOOL_ERROR"})
     void recordsBusinessOutcomeNotOnlyHttpStatus(final String mode, final 
Outcome expected) {
diff --git 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/remote/RequestScopedMcpClientTest.java
 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/remote/RequestScopedMcpClientTest.java
index ca83a72828..e76052a363 100644
--- 
a/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/remote/RequestScopedMcpClientTest.java
+++ 
b/shenyu-plugin/shenyu-plugin-agent-gateway/src/test/java/org/apache/shenyu/plugin/agent/gateway/remote/RequestScopedMcpClientTest.java
@@ -22,6 +22,19 @@ import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.sun.net.httpserver.HttpServer;
 import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.spec.McpSchema;
+import org.apache.shenyu.common.constant.Constants;
+import org.apache.shenyu.common.dto.AgentGatewayMcpConfig;
+import org.apache.shenyu.plugin.agent.gateway.AgentTrafficContext;
+import org.apache.shenyu.common.metrics.AgentMcpCallObserver.Outcome;
+import org.apache.shenyu.common.metrics.AgentMcpCallObserver;
+import org.apache.shenyu.plugin.agent.gateway.protocol.AgentMcpRemoteCatalog;
+import org.apache.shenyu.plugin.agent.gateway.protocol.AgentMcpDispatcher;
+import 
org.apache.shenyu.plugin.agent.gateway.protocol.AgentMcpExecutionContext;
+import org.apache.shenyu.plugin.agent.gateway.protocol.AgentMcpHttpHandler;
+import org.apache.shenyu.plugin.agent.gateway.security.AgentMcpIdentity;
+import org.apache.shenyu.plugin.agent.gateway.tool.AgentToolRegistry;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -30,6 +43,9 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
+import org.springframework.http.MediaType;
+import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
+import org.springframework.mock.web.server.MockServerWebExchange;
 
 import java.net.InetSocketAddress;
 import java.net.URI;
@@ -53,6 +69,7 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 
@@ -65,6 +82,109 @@ class RequestScopedMcpClientTest {
 
     private static final Duration WAIT = Duration.ofSeconds(5);
 
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    void expiredDeadlineIsATimeoutWithoutAnUpstreamRequest(final boolean sse) 
throws Exception {
+        List<Map<String, Object>> events = new CopyOnWriteArrayList<>();
+        try (Fixture fixture = new Fixture(sse)) {
+            RequestScopedMcpClient client = fixture.client(events::add);
+            try {
+                client.initialize().block(WAIT);
+                int before = events.size();
+                int requestsBefore = fixture.requests.size();
+                StepVerifier.create(client.callRaw(new 
McpSchema.CallToolRequest("lookup", Map.of("token", "expired")))
+                                .contextWrite(context -> 
context.put(McpTransportContext.KEY, 
McpTransportContext.create(Map.of("deadline", Instant.EPOCH)))))
+                        .expectError(TimeoutException.class).verify(WAIT);
+                assertEquals(before, events.size());
+                assertEquals(requestsBefore, fixture.requests.size());
+                assertEquals(0, client.pending());
+            } finally {
+                client.closeGracefully().block(WAIT);
+            }
+        }
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    void realRequestDeadlineReleasesItsFutureAndKeepsPeerHealthy(final boolean 
sse) throws Exception {
+        try (Fixture fixture = new Fixture(sse)) {
+            RequestScopedMcpClient client = fixture.client();
+            try {
+                client.initialize().block(WAIT);
+                call(client, "warm-deadline", 0, "a").block(WAIT);
+                StepVerifier.create(client.callRaw(new 
McpSchema.CallToolRequest("lookup", Map.of("token", "deadline", "delay", 3000)))
+                                .contextWrite(context -> 
context.put(McpTransportContext.KEY,
+                                        
McpTransportContext.create(Map.of("deadline", Instant.now().plusSeconds(1))))))
+                        .expectErrorMatches(error -> error instanceof 
TimeoutException || error instanceof 
java.net.http.HttpTimeoutException).verify(WAIT);
+                awaitIdle(client);
+                assertEquals("peer", call(client, "peer", 0, 
"b").block(WAIT).path("structuredContent").path("token").textValue());
+                assertEquals(3, fixture.requests.stream().filter(request -> 
"tools/call".equals(request.get("method"))).count());
+            } finally {
+                client.closeGracefully().block(WAIT);
+            }
+        }
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    void realRemoteDeadlineHasStableResponseMetricAndHealthyPeer(final boolean 
sse) throws Exception {
+        try (Fixture fixture = new Fixture(sse)) {
+            RequestScopedMcpClient client = fixture.client();
+            try {
+                client.initialize().block(WAIT);
+                call(client, "warm-handler", 0, "a").block(WAIT);
+                AgentMcpRemoteCatalog.Snapshot snapshot = new 
AgentMcpRemoteCatalog.Snapshot() {
+                    @Override
+                    public Map<String, ObjectNode> definitions() {
+                        ObjectNode tool = new 
ObjectMapper().createObjectNode().put("name", 
"orders.lookup").put("description", "Read fixture");
+                        tool.putObject("inputSchema").put("type", "object");
+                        return Map.of("orders.lookup", tool);
+                    }
+
+                    @Override
+                    public Mono<ObjectNode> invoke(final String name, final 
JsonObject arguments, final AgentMcpExecutionContext context) {
+                        return client.callRaw(new 
McpSchema.CallToolRequest("lookup", Map.of("token", 
arguments.get("token").getAsString(),
+                                        "delay", 
arguments.get("delay").getAsInt())))
+                                .contextWrite(value -> 
value.put(McpTransportContext.KEY, 
McpTransportContext.create(Map.of("deadline", context.getDeadline()))));
+                    }
+                };
+                AgentMcpRemoteCatalog catalog = operation -> 
operation.apply(snapshot);
+                final AgentMcpHttpHandler handler = new 
AgentMcpHttpHandler(new AgentMcpDispatcher(new AgentToolRegistry(List.of()), 
"fixture", "1", catalog),
+                        ignored -> Mono.just(new 
AgentMcpIdentity("trusted-fixture", Set.of("orders.lookup"))));
+                final AgentGatewayMcpConfig config = 
AgentGatewayMcpConfig.parse(JsonParser.parseString("{\"allowedTools\":[\"orders.lookup\"],\"timeoutMs\":3000,"
+                        + "\"responseMode\":\"" + (sse ? "sse" : "json") + 
"\"}").getAsJsonObject(), true);
+                final AgentTrafficContext traffic = new 
AgentTrafficContext("private-request", "mcp", "selector", "rule");
+                List<Outcome> outcomes = new CopyOnWriteArrayList<>();
+                AgentMcpCallObserver observer = (outcome, millis) -> 
outcomes.add(outcome);
+                MockServerWebExchange timedOut = 
gatewayRequest("slow-handler", 6000, observer);
+                StepVerifier.create(handler.handle(timedOut, config, traffic, 
1)).verifyComplete();
+                assertEquals(504, 
timedOut.getResponse().getStatusCode().value());
+                assertEquals("{\"error\":\"Request deadline exceeded\"}", 
timedOut.getResponse().getBodyAsString().block(WAIT));
+                awaitIdle(client);
+                MockServerWebExchange peer = gatewayRequest("peer", 0, 
observer);
+                StepVerifier.create(handler.handle(peer, config, traffic, 
2)).verifyComplete();
+                assertEquals(200, peer.getResponse().getStatusCode().value());
+                
assertTrue(peer.getResponse().getBodyAsString().block(WAIT).contains("peer"));
+                assertEquals(List.of(Outcome.TIMEOUT, Outcome.TOOL_ERROR), 
outcomes);
+                assertEquals(3, fixture.requests.stream().filter(request -> 
"tools/call".equals(request.get("method"))).count());
+                awaitIdle(client);
+            } finally {
+                client.closeGracefully().block(WAIT);
+            }
+        }
+    }
+
+    private MockServerWebExchange gatewayRequest(final String token, final int 
delay, final AgentMcpCallObserver observer) {
+        String body = 
"{\"jsonrpc\":\"2.0\",\"id\":\"same-id\",\"method\":\"tools/call\",\"params\":{\"name\":\"orders.lookup\","
+                + "\"arguments\":{\"token\":\"" + token + "\",\"delay\":" + 
delay + 
"},\"_meta\":{\"io.modelcontextprotocol/protocolVersion\":\"2026-07-28\","
+                + "\"io.modelcontextprotocol/clientCapabilities\":{}}}}";
+        MockServerWebExchange exchange = 
MockServerWebExchange.from(MockServerHttpRequest.post("/mcp").contentType(MediaType.APPLICATION_JSON)
+                .header("Accept", "application/json, 
text/event-stream").header("MCP-Protocol-Version", "2026-07-28")
+                .header("Mcp-Method", "tools/call").header("Mcp-Name", 
"orders.lookup").body(body));
+        exchange.getAttributes().put(Constants.METRICS_AGENT_MCP_CALL, 
observer);
+        return exchange;
+    }
+
     @ParameterizedTest
     @ValueSource(booleans = {false, true})
     void responseDiagnosticsHaveFixedShapeAndNoRequestData(final boolean sse) 
throws Exception {

Reply via email to