diff --git a/src/main/java/com/google/cloud/mcp/McpToolboxClient.java b/src/main/java/com/google/cloud/mcp/McpToolboxClient.java index ff79e04..88f6dcf 100644 --- a/src/main/java/com/google/cloud/mcp/McpToolboxClient.java +++ b/src/main/java/com/google/cloud/mcp/McpToolboxClient.java @@ -175,6 +175,22 @@ interface Builder { */ Builder protocolVersion(ProtocolVersion protocolVersion); + /** + * Sets the connect timeout for the underlying HttpClient. + * + * @param connectTimeout The connect timeout. + * @return The builder instance. + */ + Builder connectTimeout(java.time.Duration connectTimeout); + + /** + * Sets the request timeout for every HTTP request. + * + * @param requestTimeout The request timeout. + * @return The builder instance. + */ + Builder requestTimeout(java.time.Duration requestTimeout); + /** * Sets a custom {@link java.net.http.HttpClient} for connection management. * @@ -191,6 +207,14 @@ interface Builder { */ Builder executor(java.util.concurrent.Executor executor); + /** + * Sets a custom Logger for telemetry and logs. + * + * @param logger The custom Logger. + * @return The builder instance. + */ + Builder logger(java.util.logging.Logger logger); + /** * Builds and returns a new {@link McpToolboxClient} instance. * diff --git a/src/main/java/com/google/cloud/mcp/client/McpToolboxClientBuilder.java b/src/main/java/com/google/cloud/mcp/client/McpToolboxClientBuilder.java index 0b66c13..7efc8dc 100644 --- a/src/main/java/com/google/cloud/mcp/client/McpToolboxClientBuilder.java +++ b/src/main/java/com/google/cloud/mcp/client/McpToolboxClientBuilder.java @@ -38,8 +38,11 @@ public final class McpToolboxClientBuilder implements McpToolboxClient.Builder { private final List preProcessors = new ArrayList<>(); private final List postProcessors = new ArrayList<>(); private ProtocolVersion protocolVersion; + private java.time.Duration connectTimeout; + private java.time.Duration requestTimeout; private java.net.http.HttpClient httpClient; private java.util.concurrent.Executor executor; + private java.util.logging.Logger logger; /** Constructs a new McpToolboxClientBuilder. */ public McpToolboxClientBuilder() {} @@ -92,6 +95,18 @@ public McpToolboxClient.Builder protocolVersion(ProtocolVersion protocolVersion) return this; } + @Override + public McpToolboxClient.Builder connectTimeout(java.time.Duration connectTimeout) { + this.connectTimeout = connectTimeout; + return this; + } + + @Override + public McpToolboxClient.Builder requestTimeout(java.time.Duration requestTimeout) { + this.requestTimeout = requestTimeout; + return this; + } + @Override public McpToolboxClient.Builder httpClient(java.net.http.HttpClient httpClient) { this.httpClient = httpClient; @@ -104,6 +119,12 @@ public McpToolboxClient.Builder executor(java.util.concurrent.Executor executor) return this; } + @Override + public McpToolboxClient.Builder logger(java.util.logging.Logger logger) { + this.logger = logger; + return this; + } + @Override public McpToolboxClient build() { if (baseUrl == null || baseUrl.isEmpty()) { @@ -137,7 +158,10 @@ public McpToolboxClient build() { resolvedProvider, this.protocolVersion, this.httpClient, - this.executor); + this.executor, + this.connectTimeout, + this.requestTimeout, + this.logger); return new McpToolboxClientImpl( transport, this.headers, resolvedProvider, preProcessors, postProcessors); } diff --git a/src/main/java/com/google/cloud/mcp/transport/BaseMcpTransport.java b/src/main/java/com/google/cloud/mcp/transport/BaseMcpTransport.java index a09b383..818b1a5 100644 --- a/src/main/java/com/google/cloud/mcp/transport/BaseMcpTransport.java +++ b/src/main/java/com/google/cloud/mcp/transport/BaseMcpTransport.java @@ -36,22 +36,60 @@ import java.util.concurrent.CompletableFuture; import java.util.logging.Logger; +/** + * Base class for HTTP-based MCP transports providing common functionality like session tracking, + * header merging, and credentials resolution. + */ public abstract class BaseMcpTransport implements Transport { + /** Default static logger for BaseMcpTransport. */ protected static final Logger logger = Logger.getLogger(BaseMcpTransport.class.getName()); + + /** Warning message displayed when using unencrypted HTTP connections. */ protected static final String HTTP_WARNING = "This connection is using HTTP. To prevent credential exposure, please ensure all" + " communication is sent over HTTPS."; + /** The base URL of the MCP service. */ protected final String baseUrl; + + /** Client headers configured for the transport. */ protected final Map clientHeaders; + + /** The credentials provider for dynamic authorization. */ protected final CredentialsProvider credentialsProvider; + + /** The HTTP client used for requests. */ protected final HttpClient httpClient; + + /** The ObjectMapper for JSON serialization. */ protected final ObjectMapper objectMapper; + + /** The preferred protocol version. */ protected final ProtocolVersion preferredProtocolVersion; + + /** Lock object to synchronize initialization. */ protected final Object initLock = new Object(); + + /** Future indicating the status of initialization. */ protected CompletableFuture initFuture; + /** The request timeout for HTTP requests. */ + protected final java.time.Duration requestTimeout; + + /** The logger used for active transport logging. */ + protected final Logger activeLogger; + + /** + * Constructs a new BaseMcpTransport. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param preferredProtocolVersion The preferred protocol version. + * @param httpClient The HTTP client. + * @param executor The executor. + */ protected BaseMcpTransport( final String baseUrl, final Map clientHeaders, @@ -59,6 +97,41 @@ protected BaseMcpTransport( final ProtocolVersion preferredProtocolVersion, final HttpClient httpClient, final java.util.concurrent.Executor executor) { + this( + baseUrl, + clientHeaders, + credentialsProvider, + preferredProtocolVersion, + httpClient, + executor, + null, + null, + null); + } + + /** + * Constructs a new BaseMcpTransport with timeouts and custom logger. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param preferredProtocolVersion The preferred protocol version. + * @param httpClient The HTTP client. + * @param executor The executor. + * @param connectTimeout The connection timeout. + * @param requestTimeout The request timeout. + * @param logger The custom logger. + */ + protected BaseMcpTransport( + final String baseUrl, + final Map clientHeaders, + final CredentialsProvider credentialsProvider, + final ProtocolVersion preferredProtocolVersion, + final HttpClient httpClient, + final java.util.concurrent.Executor executor, + final java.time.Duration connectTimeout, + final java.time.Duration requestTimeout, + final Logger logger) { if (baseUrl == null || baseUrl.isEmpty()) { throw new IllegalArgumentException("Base URL must be provided"); } @@ -78,12 +151,14 @@ protected BaseMcpTransport( HttpClient.Builder builder = HttpClient.newBuilder() .cookieHandler(new java.net.CookieManager()) - .connectTimeout(Duration.ofSeconds(10)); + .connectTimeout(connectTimeout != null ? connectTimeout : Duration.ofSeconds(10)); if (executor != null) { builder.executor(executor); } this.httpClient = builder.build(); } + this.requestTimeout = requestTimeout; + this.activeLogger = logger != null ? logger : BaseMcpTransport.logger; this.objectMapper = new ObjectMapper(); } @@ -186,9 +261,21 @@ final CompletableFuture ensureInitialized(final Map extraM } } + /** + * Performs the version-specific initialization handshake. + * + * @param authHeader The authorization header value, if present. + * @param handshakeHeaders The resolved headers for the handshake. + * @return A CompletableFuture that completes when initialization is done. + */ protected abstract CompletableFuture performInitialization( final String authHeader, final Map handshakeHeaders); + /** + * Applies protocol-specific headers to the request builder. + * + * @param builder The HTTP request builder. + */ protected abstract void applyProtocolHeaders(final HttpRequest.Builder builder); @Override @@ -196,7 +283,7 @@ public final CompletableFuture listTools( final String toolsetName, final Map metadata) { if (this.baseUrl.toLowerCase(java.util.Locale.ROOT).startsWith("http://") && !metadata.isEmpty()) { - logger.warning(HTTP_WARNING); + activeLogger.warning(HTTP_WARNING); } return ensureInitialized(metadata) .thenCompose(v -> mergeHeaders(metadata)) @@ -211,6 +298,9 @@ public final CompletableFuture listTools( HttpRequest.newBuilder() .uri(URI.create(url)) .POST(HttpRequest.BodyPublishers.ofString(body)); + if (requestTimeout != null) { + req.timeout(requestTimeout); + } mergedHeaders.forEach(req::setHeader); applyProtocolHeaders(req); @@ -230,7 +320,7 @@ public final CompletableFuture invokeTool( final Map metadata) { if (this.baseUrl.toLowerCase(java.util.Locale.ROOT).startsWith("http://") && !metadata.isEmpty()) { - logger.warning(HTTP_WARNING); + activeLogger.warning(HTTP_WARNING); } return ensureInitialized(metadata) .thenCompose(v -> mergeHeaders(metadata)) @@ -247,6 +337,9 @@ public final CompletableFuture invokeTool( .uri(URI.create(baseUrl)) .POST(HttpRequest.BodyPublishers.ofString(requestBody)); + if (requestTimeout != null) { + requestBuilder.timeout(requestTimeout); + } mergedHeaders.forEach(requestBuilder::setHeader); applyProtocolHeaders(requestBuilder); diff --git a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransport.java b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransport.java index b034a87..863e6ed 100644 --- a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransport.java +++ b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransport.java @@ -19,8 +19,10 @@ import com.google.cloud.mcp.ProtocolVersion; import com.google.cloud.mcp.auth.CredentialsProvider; import java.net.http.HttpClient; +import java.time.Duration; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.logging.Logger; /** Default HTTP transport implementation routing requests to version-specific handlers. */ public final class HttpMcpTransport implements Transport { @@ -95,6 +97,41 @@ public HttpMcpTransport( final ProtocolVersion preferredProtocolVersion, final HttpClient httpClient, final java.util.concurrent.Executor executor) { + this( + baseUrl, + clientHeaders, + credentialsProvider, + preferredProtocolVersion, + httpClient, + executor, + null, + null, + null); + } + + /** + * Constructs a new HttpMcpTransport with full configuration. + * + * @param baseUrl The base URL of the MCP service. + * @param clientHeaders Optional headers to include in every request. + * @param credentialsProvider Optional provider for auth credentials. + * @param preferredProtocolVersion Optional preferred protocol version. + * @param httpClient Optional HttpClient instance. + * @param executor Optional Executor for handling async requests. + * @param connectTimeout Optional connection timeout. + * @param requestTimeout Optional request timeout. + * @param logger Optional Logger instance. + */ + public HttpMcpTransport( + final String baseUrl, + final Map clientHeaders, + final CredentialsProvider credentialsProvider, + final ProtocolVersion preferredProtocolVersion, + final HttpClient httpClient, + final java.util.concurrent.Executor executor, + final Duration connectTimeout, + final Duration requestTimeout, + final Logger logger) { final ProtocolVersion version = preferredProtocolVersion != null ? preferredProtocolVersion @@ -104,34 +141,73 @@ public HttpMcpTransport( case VERSION_2025_11_25: this.delegate = new HttpMcpTransportV20251125( - baseUrl, clientHeaders, credentialsProvider, httpClient, executor); + baseUrl, + clientHeaders, + credentialsProvider, + httpClient, + executor, + connectTimeout, + requestTimeout, + logger); break; case VERSION_2025_06_18: this.delegate = new HttpMcpTransportV20250618( - baseUrl, clientHeaders, credentialsProvider, httpClient, executor); + baseUrl, + clientHeaders, + credentialsProvider, + httpClient, + executor, + connectTimeout, + requestTimeout, + logger); break; case VERSION_2025_03_26: this.delegate = new HttpMcpTransportV20250326( - baseUrl, clientHeaders, credentialsProvider, httpClient, executor); + baseUrl, + clientHeaders, + credentialsProvider, + httpClient, + executor, + connectTimeout, + requestTimeout, + logger); break; case VERSION_2024_11_05: this.delegate = new HttpMcpTransportV20241105( - baseUrl, clientHeaders, credentialsProvider, httpClient, executor); + baseUrl, + clientHeaders, + credentialsProvider, + httpClient, + executor, + connectTimeout, + requestTimeout, + logger); break; default: throw new IllegalArgumentException("Unsupported protocol version: " + version); } } - /** Internal constructor for testing purposes. */ + /** + * Internal constructor for testing purposes. + * + * @param baseUrl The base URL. + * @param httpClient The mock HttpClient. + */ public HttpMcpTransport(final String baseUrl, final HttpClient httpClient) { this(baseUrl, Map.of(), null, null, httpClient, null); } - /** Internal constructor for testing purposes. */ + /** + * Internal constructor for testing purposes. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param httpClient The mock HttpClient. + */ public HttpMcpTransport( final String baseUrl, final Map clientHeaders, final HttpClient httpClient) { this(baseUrl, clientHeaders, null, null, httpClient, null); diff --git a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20241105.java b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20241105.java index f9a0576..9a7e39a 100644 --- a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20241105.java +++ b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20241105.java @@ -24,11 +24,23 @@ import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.time.Duration; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.logging.Logger; +/** HTTP transport implementation for protocol version 2024-11-05. */ public final class HttpMcpTransportV20241105 extends BaseMcpTransport { + /** + * Constructs a new HttpMcpTransportV20241105. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param httpClient The HTTP client. + * @param executor The executor. + */ public HttpMcpTransportV20241105( final String baseUrl, final Map clientHeaders, @@ -44,6 +56,39 @@ public HttpMcpTransportV20241105( executor); } + /** + * Constructs a new HttpMcpTransportV20241105 with timeouts and logger. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param httpClient The HTTP client. + * @param executor The executor. + * @param connectTimeout The connection timeout. + * @param requestTimeout The request timeout. + * @param logger The logger. + */ + public HttpMcpTransportV20241105( + final String baseUrl, + final Map clientHeaders, + final CredentialsProvider credentialsProvider, + final HttpClient httpClient, + final java.util.concurrent.Executor executor, + final Duration connectTimeout, + final Duration requestTimeout, + final Logger logger) { + super( + baseUrl, + clientHeaders, + credentialsProvider, + ProtocolVersion.VERSION_2024_11_05, + httpClient, + executor, + connectTimeout, + requestTimeout, + logger); + } + @Override protected CompletableFuture performInitialization( final String authHeader, final Map handshakeHeaders) { diff --git a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20250326.java b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20250326.java index d642076..1ebafd3 100644 --- a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20250326.java +++ b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20250326.java @@ -24,14 +24,26 @@ import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.time.Duration; import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.logging.Logger; +/** HTTP transport implementation for protocol version 2025-03-26. */ public final class HttpMcpTransportV20250326 extends BaseMcpTransport { private volatile String sessionId; + /** + * Constructs a new HttpMcpTransportV20250326. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param httpClient The HTTP client. + * @param executor The executor. + */ public HttpMcpTransportV20250326( final String baseUrl, final Map clientHeaders, @@ -47,6 +59,39 @@ public HttpMcpTransportV20250326( executor); } + /** + * Constructs a new HttpMcpTransportV20250326 with timeouts and logger. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param httpClient The HTTP client. + * @param executor The executor. + * @param connectTimeout The connection timeout. + * @param requestTimeout The request timeout. + * @param logger The logger. + */ + public HttpMcpTransportV20250326( + final String baseUrl, + final Map clientHeaders, + final CredentialsProvider credentialsProvider, + final HttpClient httpClient, + final java.util.concurrent.Executor executor, + final Duration connectTimeout, + final Duration requestTimeout, + final Logger logger) { + super( + baseUrl, + clientHeaders, + credentialsProvider, + ProtocolVersion.VERSION_2025_03_26, + httpClient, + executor, + connectTimeout, + requestTimeout, + logger); + } + @Override protected CompletableFuture performInitialization( final String authHeader, final Map handshakeHeaders) { diff --git a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20250618.java b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20250618.java index 8ac79c9..e444fbc 100644 --- a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20250618.java +++ b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20250618.java @@ -24,11 +24,23 @@ import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.time.Duration; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.logging.Logger; +/** HTTP transport implementation for protocol version 2025-06-18. */ public final class HttpMcpTransportV20250618 extends BaseMcpTransport { + /** + * Constructs a new HttpMcpTransportV20250618. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param httpClient The HTTP client. + * @param executor The executor. + */ public HttpMcpTransportV20250618( final String baseUrl, final Map clientHeaders, @@ -44,6 +56,39 @@ public HttpMcpTransportV20250618( executor); } + /** + * Constructs a new HttpMcpTransportV20250618 with timeouts and logger. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param httpClient The HTTP client. + * @param executor The executor. + * @param connectTimeout The connection timeout. + * @param requestTimeout The request timeout. + * @param logger The logger. + */ + public HttpMcpTransportV20250618( + final String baseUrl, + final Map clientHeaders, + final CredentialsProvider credentialsProvider, + final HttpClient httpClient, + final java.util.concurrent.Executor executor, + final Duration connectTimeout, + final Duration requestTimeout, + final Logger logger) { + super( + baseUrl, + clientHeaders, + credentialsProvider, + ProtocolVersion.VERSION_2025_06_18, + httpClient, + executor, + connectTimeout, + requestTimeout, + logger); + } + @Override protected CompletableFuture performInitialization( final String authHeader, final Map handshakeHeaders) { diff --git a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20251125.java b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20251125.java index 4ab7f37..40f9149 100644 --- a/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20251125.java +++ b/src/main/java/com/google/cloud/mcp/transport/HttpMcpTransportV20251125.java @@ -24,11 +24,23 @@ import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.time.Duration; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.logging.Logger; +/** HTTP transport implementation for protocol version 2025-11-25. */ public final class HttpMcpTransportV20251125 extends BaseMcpTransport { + /** + * Constructs a new HttpMcpTransportV20251125. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param httpClient The HTTP client. + * @param executor The executor. + */ public HttpMcpTransportV20251125( final String baseUrl, final Map clientHeaders, @@ -44,6 +56,39 @@ public HttpMcpTransportV20251125( executor); } + /** + * Constructs a new HttpMcpTransportV20251125 with timeouts and logger. + * + * @param baseUrl The base URL. + * @param clientHeaders The client headers. + * @param credentialsProvider The credentials provider. + * @param httpClient The HTTP client. + * @param executor The executor. + * @param connectTimeout The connection timeout. + * @param requestTimeout The request timeout. + * @param logger The logger. + */ + public HttpMcpTransportV20251125( + final String baseUrl, + final Map clientHeaders, + final CredentialsProvider credentialsProvider, + final HttpClient httpClient, + final java.util.concurrent.Executor executor, + final Duration connectTimeout, + final Duration requestTimeout, + final Logger logger) { + super( + baseUrl, + clientHeaders, + credentialsProvider, + ProtocolVersion.VERSION_2025_11_25, + httpClient, + executor, + connectTimeout, + requestTimeout, + logger); + } + @Override protected CompletableFuture performInitialization( final String authHeader, final Map handshakeHeaders) { diff --git a/src/test/java/com/google/cloud/mcp/transport/HttpMcpTransportTest.java b/src/test/java/com/google/cloud/mcp/transport/HttpMcpTransportTest.java index 0953558..7921c15 100644 --- a/src/test/java/com/google/cloud/mcp/transport/HttpMcpTransportTest.java +++ b/src/test/java/com/google/cloud/mcp/transport/HttpMcpTransportTest.java @@ -33,6 +33,7 @@ import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.time.Duration; import java.util.Collections; import java.util.List; import java.util.Map; @@ -280,7 +281,6 @@ void testConstructor_WithCustomExecutorConfiguresHttpClient() throws Exception { } } - @Test @SuppressWarnings("unchecked") void testInitialize_ServerReturnsErrorJsonRpcResponse() throws Exception { HttpResponse mockInitResponse = mock(HttpResponse.class); @@ -476,9 +476,47 @@ void testListTools_ParsesComplexToolsCorrectly() throws Exception { } @Test - void testJsonRpcInstantiation() { - // Instantiate package-private JsonRpc namespace to cover its default constructor - JsonRpc rpc = new JsonRpc(); - assertNotNull(rpc); + @SuppressWarnings("unchecked") + void testInvokeTool_WithRequestTimeout() throws Exception { + HttpMcpTransport transportWithTimeout = + new HttpMcpTransport( + "https://test-mcp-service.com", + Map.of(), + null, + ProtocolVersion.VERSION_2025_11_25, + mockClient, + null, + Duration.ofSeconds(5), + Duration.ofSeconds(3), + null); + + HttpResponse mockInitResponse = mock(HttpResponse.class); + when(mockInitResponse.statusCode()).thenReturn(200); + when(mockInitResponse.body()) + .thenReturn( + "{\"jsonrpc\":\"2.0\",\"id\":\"1\",\"result\":{\"protocolVersion\":\"2025-11-25\"}}"); + + HttpResponse mockInitializedResponse = mock(HttpResponse.class); + when(mockInitializedResponse.statusCode()).thenReturn(200); + when(mockInitializedResponse.body()).thenReturn(""); + + HttpResponse mockInvokeResponse = mock(HttpResponse.class); + when(mockInvokeResponse.statusCode()).thenReturn(200); + when(mockInvokeResponse.body()) + .thenReturn( + "{\"jsonrpc\":\"2.0\",\"id\":\"3\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"success\"}]}}"); + + when(mockClient.sendAsync(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))) + .thenReturn(CompletableFuture.completedFuture(mockInitResponse)) + .thenReturn(CompletableFuture.completedFuture(mockInitializedResponse)) + .thenReturn(CompletableFuture.completedFuture(mockInvokeResponse)); + + TransportResponse response = + transportWithTimeout + .invokeTool("test-tool", Map.of("param1", "value1"), Collections.emptyMap()) + .get(); + + assertNotNull(response); + assertEquals(200, response.getStatusCode()); } }