diff --git a/.changeset/cap-agent-response-size.md b/.changeset/cap-agent-response-size.md new file mode 100644 index 0000000..1f0bcca --- /dev/null +++ b/.changeset/cap-agent-response-size.md @@ -0,0 +1,56 @@ +--- +"adcp": minor +--- + +feat(transport): cap agent response size at 4 MiB by default + +Agent responses are untrusted input. Every MCP and A2A call response — tool calls, `tools/list`, +SSE streams and agent-card fetches — is now capped at **4 MiB** by default, and the cap is enforced +while the body streams in: a declared `Content-Length` over the cap is rejected before a byte is +read, and a chunked body is aborted the moment the running count crosses it. SSE is capped per +event as well as cumulatively per call, so a peer cannot buy extra budget by answering with +`Content-Type: text/event-stream`. Outbound requests force `Accept-Encoding: identity` so the +counted bytes are the bytes you hold. + +Previously the only limit on the call paths was a 10 MB *character* check that ran after the whole +body had been buffered and parsed — far too late to stop a memory exhaustion. One agent answering +`get_products` with a 17.8 MB payload was enough to OOM-kill a hosted grader at 10+ GB RSS. + +**This is a behaviour change.** A call that used to return a 6 MB catalogue now throws +`ResponseTooLargeError` (`code() == "RESPONSE_TOO_LARGE"`). Raise the limit where you need it — +there is no opt-out, only a bigger number: + +```java +// per client +AdcpClient.builder().agent(agent).maxResponseBytes(32L * 1024 * 1024).build(); + +// per agent +AgentConfig.builder().id("big-catalogue").agentUri(uri) + .maxResponseBytes(32L * 1024 * 1024).build(); + +// per call +client.callTool("get_products", args, GetProductsResponse.class, + CallToolOptions.builder().maxResponseBytes(32L * 1024 * 1024).build()); +``` + +The cap is enforced inside a JDK body subscriber and may reach you wrapped, so branch on it with +`ResponseTooLargeError.findIn(throwable)`. + +Also in this change: + +- `CallToolOptions.maxResponseBytes` is now enforced (it was documented as "reserved, not yet + enforced"). `CallToolOptions.timeout` remains reserved. +- `AgentConfig` gains a `maxResponseBytes` component, and `AdcpClient.Builder` / + `AgentConfig.Builder` gain `maxResponseBytes(long)`. Code that calls the `AgentConfig` canonical + constructor positionally needs one more argument; the builder is unaffected. +- Agent-card fetches moved from the 4 KiB discovery-probe cap to the 4 MiB call cap and now throw + instead of silently truncating — a card over 4 KiB used to be swapped for a synthetic fallback. +- MCP and A2A connection caches are keyed by the response cap, so a call that overrides it does not + reuse a connection built for a different one. +- Agent-card fetches get their own read deadline. `HttpRequest.timeout` only covers status line and + headers, so a peer that sends headers and then drips or stalls could otherwise pin a thread for as + long as it liked. + +Defaults, overrides and the SSE rules are documented in `docs/response-size-limits.md`, with the +design in `specs/response-size-limits.md`. +Cross-SDK contract: https://github.com/adcontextprotocol/adcp/issues/8190 diff --git a/ROADMAP.md b/ROADMAP.md index c1e57e0..cb34824 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -33,6 +33,7 @@ Decisions made post-RFC that supersede or refine the merged text. The table belo | D21 | Branch protection on `main` | **Required: 1 code-owner approving review, required CI checks (`build`, `test`, `storyboard`, `IPR Policy / Signature`), no force-push, no direct push, no admin bypass.** Dependabot patch PRs auto-merge after green CI. Two-reviewer gate doesn't add safety with a founder pair; revisit when contributor base grows. | (Sets a repo convention) | | D22 | Multi-tenant signing keys | **Signing key selection is tenant-aware, not just `adcp_use`-aware.** The L1 signing SPI must select by purpose plus caller/resolved publisher tenant context, a JOSE-compatible pattern where one operator can publish multiple tenant keys on a single JWKS endpoint: same `adcp_use`, distinct `kid` per tenant. v0.2 freezes the explicit parameter set for `SigningContext` — `AdcpUse use()`, nullable tenant identity, nullable principal reference — and the shipped API is `SigningContext`-based with no single-arg `SigningProvider.forUse(AdcpUse)` short form. v0.3 wires tenant resolution from `AccountStore` / `adagents.json` into that selector. Longer-form shape lives in [`specs/signing-context.md`](specs/signing-context.md). | Refines RFC §L1 signing and §L2 multi-tenant principal resolution | | D23 | AdCP 3.1 beta parity target | **Track the signed protocol bundle as the source of truth; use TS/Python as implementation references.** For this roadmap update the pinned target is `3.1.0-beta.5` (2026-05-26). If 3.1 GA cuts before the v0.1 codegen PR, bump to GA deliberately in that PR; otherwise beta.5 is the target. `@adcp/sdk@8.1.0-beta.13` is aligned to beta.5; Python `6.3.0b4` is on beta.4 and is a reference, not a blocker. Java's local `ADCP_VERSION` remains `3.0.11`; v0.1 must include wire-version negotiation, multi-bundle schema loading, and 3.1 beta/GA codegen coverage before claiming cross-SDK parity. | Updates RFC parity baseline and D15 spec-rev SLO | +| D24 | Agent response size cap | **Agent responses are capped by default at 4 MiB, enforced while the body streams in, throwing `ResponseTooLargeError`.** Covers MCP `callTool`, A2A `message/send`, SSE (per event and cumulatively per call), `tools/list` and agent-card fetches. Overridable per client, per agent and per call; never disablable. Enforced by wrapping the JDK `HttpClient` handed to the third-party MCP/A2A transports, which expose no body-size hook of their own. Discovery probes keep their separate 4 KiB cap from [`specs/ssrf-baseline.md`](specs/ssrf-baseline.md). Full shape in [`specs/response-size-limits.md`](specs/response-size-limits.md), adopter guide in [`docs/response-size-limits.md`](docs/response-size-limits.md); cross-SDK contract in [adcp#8190](https://github.com/adcontextprotocol/adcp/issues/8190). | Extends RFC §SSRF baseline body cap from probes to call paths | ## Parity baseline (as of 2026-05-26) @@ -587,7 +588,7 @@ Additional decisions added post-RFC that remain open: | RFC merged on adcontextprotocol/adcp | ✅ (PR #4279, 2026-05-13) | | RFC imported to this repo | ✅ ([docs/rfc/java-sdk-rfc.md](docs/rfc/java-sdk-rfc.md)) | | Implementation plan drafted | ✅ (this doc) | -| Confirmed decisions | ✅ D1–D21 locked; D22 + D23 framing locked, shape/value details captured in this roadmap and linked specs. | +| Confirmed decisions | ✅ D1–D21 locked; D22 + D23 framing locked, D24 shipped, shape/value details captured in this roadmap and linked specs. | | Protocol parity check | 🟡 Behind — local Java pin is `3.0.11`; latest upstream check is AdCP `3.1.0-beta.5`, TS `@adcp/sdk@8.1.0-beta.13`, Python `adcp@6.3.0b4` on `3.1.0-beta.4`. | | Funding / staffing confirmed | ⏳ Decision pending | | Tracks claimed | 3 / 14 — `infra` (Track 1, #2), `codegen` (Track 2, #11), `transport` (Track 3, #17) | diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/AdcpClient.java b/adcp/src/main/java/org/adcontextprotocol/adcp/AdcpClient.java index 02e52df..73aff60 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/AdcpClient.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/AdcpClient.java @@ -35,6 +35,15 @@ *

Named convenience methods (e.g. {@code getProducts()}) are provided * for each tool in the TS SDK parity target. All funnel through * {@link #callTool(String, Map, Class, CallToolOptions)}. + * + *

Response size cap. Every agent response is capped at 4 MiB by + * default and the cap is enforced while the body streams in; an over-cap + * response throws + * {@link org.adcontextprotocol.adcp.error.ResponseTooLargeError}. Raise it with + * {@link Builder#maxResponseBytes(long)} (per client), + * {@link AgentConfig.Builder#maxResponseBytes(long)} (per agent) or + * {@link CallToolOptions.Builder#maxResponseBytes(long)} (per call). See + * {@code specs/response-size-limits.md}. */ public final class AdcpClient implements AutoCloseable { @@ -63,6 +72,7 @@ private AdcpClient(Builder builder) { this.adcpHttpClient = AdcpHttpClient.builder() .ssrfPolicy(ssrfPolicy) + .maxCallResponseBytes(builder.maxResponseBytes) .build(); McpConnectionManager mcpConnectionManager = new McpConnectionManager( Duration.ofSeconds(10), builder.requestTimeout, adcpHttpClient); @@ -174,6 +184,7 @@ public static final class Builder { private @Nullable ObjectMapper objectMapper; private @Nullable SsrfPolicy ssrfPolicy; private Duration requestTimeout = Duration.ofSeconds(30); + private long maxResponseBytes = AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES; private Builder() {} @@ -225,6 +236,34 @@ public Builder requestTimeout(Duration requestTimeout) { return this; } + /** + * Caps every agent response at {@code maxResponseBytes}. Default 4 MiB + * ({@link AdcpHttpClient#DEFAULT_MAX_CALL_RESPONSE_BYTES}). + * + *

The cap covers MCP tool calls, A2A {@code message/send}, SSE + * streams and agent-card fetches, and is enforced while the body + * streams in: an over-cap response throws + * {@link org.adcontextprotocol.adcp.error.ResponseTooLargeError} and is + * never buffered in full. Agent responses are untrusted input, so the + * cap cannot be disabled — only raised. + * + *

A single agent can override this via + * {@link AgentConfig.Builder#maxResponseBytes(long)}, and a single call + * via + * {@link org.adcontextprotocol.adcp.transport.CallToolOptions.Builder#maxResponseBytes(long)}. + * + * @param maxResponseBytes cap in bytes; must be positive + */ + public Builder maxResponseBytes(long maxResponseBytes) { + if (maxResponseBytes <= 0) { + throw new ConfigurationError( + "maxResponseBytes must be positive: " + maxResponseBytes, + "maxResponseBytes"); + } + this.maxResponseBytes = maxResponseBytes; + return this; + } + /** Builds the client. */ public AdcpClient build() { validateAdcpVersion(adcpVersion); diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/AgentConfig.java b/adcp/src/main/java/org/adcontextprotocol/adcp/AgentConfig.java index 0edd865..8bb1c03 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/AgentConfig.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/AgentConfig.java @@ -26,6 +26,10 @@ *

  • OAuth auth-code ({@code oauthTokens})
    • * * + *

    Responses from this agent are capped at + * {@link #maxResponseBytes()} bytes (4 MiB by default) and the cap is enforced + * while the body streams in — see {@code specs/response-size-limits.md}. + * * Use {@link #builder()} to construct instances. */ public record AgentConfig( @@ -39,7 +43,8 @@ public record AgentConfig( @Nullable String webhookUrlTemplate, @Nullable String webhookSecret, @Nullable AdcpVersion adcpVersion, - Map extraHeaders + Map extraHeaders, + @Nullable Long maxResponseBytes ) { private static final Logger log = LoggerFactory.getLogger(AgentConfig.class); @@ -52,6 +57,7 @@ public record AgentConfig( validateAuth(authToken, basicAuth, oauthClientCredentials, oauthTokens); validateAuthToken(authToken); validateExtraHeaders(extraHeaders); + validateMaxResponseBytes(maxResponseBytes); warnPlaintextAuth(agentUri, authToken, basicAuth, oauthClientCredentials, oauthTokens); } @@ -68,7 +74,8 @@ public String toString() { + ", webhookSecret=" + (webhookSecret != null ? "" : "null") + ", adcpVersion=" + adcpVersion + ", extraHeaders=" + (extraHeaders.isEmpty() - ? "{}" : "<" + extraHeaders.size() + " headers>") + "]"; + ? "{}" : "<" + extraHeaders.size() + " headers>") + + ", maxResponseBytes=" + maxResponseBytes + "]"; } /** Creates a builder for {@code AgentConfig}. */ @@ -137,6 +144,14 @@ private static boolean hasCrlf(String s) { return s.indexOf('\r') >= 0 || s.indexOf('\n') >= 0; } + private static void validateMaxResponseBytes(@Nullable Long maxResponseBytes) { + if (maxResponseBytes != null && maxResponseBytes <= 0) { + throw new ConfigurationError( + "maxResponseBytes must be positive: " + maxResponseBytes, + "maxResponseBytes"); + } + } + private static void warnPlaintextAuth( URI agentUri, @Nullable String authToken, @@ -166,6 +181,7 @@ public static final class Builder { private @Nullable String webhookSecret; private @Nullable AdcpVersion adcpVersion; private Map extraHeaders = Map.of(); + private @Nullable Long maxResponseBytes; private Builder() {} @@ -246,6 +262,29 @@ public Builder extraHeaders(Map extraHeaders) { return this; } + /** + * Caps this agent's responses at {@code maxResponseBytes}, overriding + * the client default (4 MiB, see + * {@link org.adcontextprotocol.adcp.http.AdcpHttpClient#DEFAULT_MAX_CALL_RESPONSE_BYTES}). + * + *

    Raise it for an agent that legitimately publishes a very large + * catalogue. There is no way to disable the cap: agent responses are + * untrusted input, and an uncapped read is how one peer memory-DoSes + * the caller. A single call can override this again via + * {@link org.adcontextprotocol.adcp.transport.CallToolOptions#maxResponseBytes()}. + * + * @param maxResponseBytes cap in bytes; must be positive + */ + public Builder maxResponseBytes(long maxResponseBytes) { + if (maxResponseBytes <= 0) { + throw new ConfigurationError( + "maxResponseBytes must be positive: " + maxResponseBytes, + "maxResponseBytes"); + } + this.maxResponseBytes = maxResponseBytes; + return this; + } + /** Builds the config, validating required fields and auth exclusivity. */ public AgentConfig build() { if (id == null) { @@ -258,7 +297,7 @@ public AgentConfig build() { id, agentUri, protocol, authToken, basicAuth, oauthClientCredentials, oauthTokens, webhookUrlTemplate, webhookSecret, adcpVersion, - extraHeaders); + extraHeaders, maxResponseBytes); } } } diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/error/ResponseTooLargeError.java b/adcp/src/main/java/org/adcontextprotocol/adcp/error/ResponseTooLargeError.java index c5441d0..0e5bf3a 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/error/ResponseTooLargeError.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/error/ResponseTooLargeError.java @@ -3,13 +3,29 @@ import org.jspecify.annotations.Nullable; import java.net.URI; +import java.net.URISyntaxException; -/** Response body exceeded the configured maximum size. */ +/** + * Response body exceeded the configured maximum size. + * + *

    Thrown by the AdCP transports when an agent answers with more bytes than + * the configured cap allows. The cap is on by default (4 MiB) and is enforced + * while the body streams in, so an oversized response never reaches the heap + * in full. See {@code specs/response-size-limits.md}. + * + *

    Raise the cap per client + * ({@code AdcpClient.builder().maxResponseBytes(...)}), per agent + * ({@code AgentConfig.builder().maxResponseBytes(...)}) or per call + * ({@code CallToolOptions.builder().maxResponseBytes(...)}). + */ public final class ResponseTooLargeError extends AdcpError { @java.io.Serial private static final long serialVersionUID = 1L; + /** Depth bound for {@link #findIn} so a self-referential cause chain cannot loop. */ + private static final int MAX_CAUSE_DEPTH = 32; + private final long limit; private final long bytesRead; private final @Nullable URI url; @@ -17,21 +33,84 @@ public final class ResponseTooLargeError extends AdcpError { public ResponseTooLargeError(long limit, long bytesRead, @Nullable URI url) { super("RESPONSE_TOO_LARGE", "Response exceeded " + limit + " bytes (read " + bytesRead + ")" - + (url != null ? " from " + url : ""), + + (url != null ? " from " + redact(url) : ""), null); this.limit = limit; this.bytesRead = bytesRead; this.url = url; } + /** + * Strips userinfo and query before the URL goes into the message. + * + *

    This message is logged, surfaced to users and often fed to an LLM, and + * agent URIs routinely carry credentials in both places. {@link #url()} + * still returns the URI the caller configured. + */ + private static URI redact(URI url) { + if (url.getRawUserInfo() == null && url.getRawQuery() == null) { + return url; + } + try { + return new URI(url.getScheme(), null, url.getHost(), url.getPort(), + url.getPath(), null, null); + } catch (URISyntaxException e) { + return URI.create(String.valueOf(url.getScheme()) + "://" + url.getHost()); + } + } + + /** + * Finds a {@code ResponseTooLargeError} in a throwable's cause chain. + * + *

    The cap is enforced inside a {@code java.net.http} body subscriber, and + * the JDK and the MCP/A2A transports each wrap subscriber failures before + * they surface. This walks the chain so callers (and the SDK's own transport + * code, which must not retry an oversized response) can recover the typed + * error. + * + * @param throwable the throwable to inspect, may be {@code null} + * @return the first {@code ResponseTooLargeError} found, or {@code null} + */ + public static @Nullable ResponseTooLargeError findIn(@Nullable Throwable throwable) { + Throwable current = throwable; + for (int depth = 0; current != null && depth < MAX_CAUSE_DEPTH; depth++) { + if (current instanceof ResponseTooLargeError error) { + return error; + } + for (Throwable suppressed : current.getSuppressed()) { + if (suppressed instanceof ResponseTooLargeError error) { + return error; + } + } + Throwable cause = current.getCause(); + if (cause == current) { + break; + } + current = cause; + } + return null; + } + + /** The configured cap, in bytes. */ public long limit() { return limit; } + /** + * Bytes counted when the cap was breached — the declared + * {@code Content-Length} for a pre-rejected response, otherwise the running + * total at the moment the read was aborted. Not the full body size: the SDK + * stops reading rather than measure an oversized payload. + */ public long bytesRead() { return bytesRead; } + /** + * The URL that produced the oversized response, if known — as configured, + * including any userinfo or query. {@link #getMessage()} carries a redacted + * form instead. + */ public @Nullable URI url() { return url; } diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/http/AdcpHttpClient.java b/adcp/src/main/java/org/adcontextprotocol/adcp/http/AdcpHttpClient.java index b9176a2..1298f64 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/http/AdcpHttpClient.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/http/AdcpHttpClient.java @@ -1,5 +1,6 @@ package org.adcontextprotocol.adcp.http; +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -12,9 +13,13 @@ import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.net.http.HttpTimeoutException; +import java.net.URISyntaxException; import java.time.Duration; import java.util.Map; import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; /** * SSRF-safe HTTP client for AdCP outbound requests. @@ -24,11 +29,38 @@ *

  • Resolve DNS once, validate the full address set
    • *
  • Validate resolved addresses against the SSRF policy
    • *
  • {@code redirect: manual} (no transparent redirect-follow)
    • - *
  • Body cap (default 4 KiB for probes, configurable per call)
    • + *
  • Body cap — two of them, for two different jobs (see below)
    • * * - *

    Every outbound HTTP call in the SDK routes through this client. - * Built on {@link java.net.http.HttpClient} (JDK 21). + *

    Two caps

    + * + * + *

    Outbound coverage. Discovery and agent-card requests are sent by + * this client directly. The MCP + * ({@code HttpClientStreamableHttpTransport} / {@code HttpClientSseClientTransport}) + * and A2A ({@code JdkA2AHttpClient}) call transports own their own sockets, so + * they do not route through {@link #send}; they are instead built on a + * {@link java.net.http.HttpClient} handed to them by + * {@link #newHttpClientBuilder()}, which carries this client's connect timeout, + * {@code followRedirects(NEVER)} policy and call-response cap. SSRF address + * validation for those paths happens in + * {@link org.adcontextprotocol.adcp.transport.ProtocolClient} before the call, + * not inside the transport. + * + *

    Built on {@link java.net.http.HttpClient} (JDK 21). * *

    The client keeps the original URI authority unchanged so HTTPS uses the * intended hostname for TLS SNI and hostname verification. DNS validation @@ -45,12 +77,27 @@ public final class AdcpHttpClient implements AutoCloseable { /** Default body cap for discovery probes (4 KiB). */ public static final long DEFAULT_MAX_RESPONSE_BYTES = 4 * 1024; + /** + * Default body cap for agent responses on the MCP and A2A call paths, + * and for agent-card fetches: 4 MiB. + * + *

    Agent responses are untrusted input. This cap is on by default so a + * single peer cannot memory-DoS the caller with an oversized payload. + * Raise it per client, per agent or per call — see + * {@code specs/response-size-limits.md}. + */ + public static final long DEFAULT_MAX_CALL_RESPONSE_BYTES = 4L * 1024 * 1024; + + /** Upper bound accepted by {@link Builder#maxResponseBytes} (64 MiB). */ + public static final long MAX_PROBE_RESPONSE_BYTES = 64L * 1024 * 1024; + private static final Duration DEFAULT_CONNECT_TIMEOUT = Duration.ofSeconds(10); private static final Duration DEFAULT_READ_TIMEOUT = Duration.ofSeconds(30); private static final String DEFAULT_USER_AGENT = "adcp-java-sdk/0.1"; private final SsrfPolicy ssrfPolicy; private final long maxResponseBytes; + private final long maxCallResponseBytes; private final Duration connectTimeout; private final Duration readTimeout; private final String userAgent; @@ -60,6 +107,7 @@ public final class AdcpHttpClient implements AutoCloseable { private AdcpHttpClient(Builder builder) { this.ssrfPolicy = builder.ssrfPolicy; this.maxResponseBytes = builder.maxResponseBytes; + this.maxCallResponseBytes = builder.maxCallResponseBytes; this.connectTimeout = builder.connectTimeout; this.readTimeout = builder.readTimeout; this.userAgent = builder.userAgent; @@ -81,8 +129,13 @@ public static Builder builder() { *

    The hostname is resolved via DNS and all addresses are validated * against the {@link SsrfPolicy}. The original URI is preserved so TLS * SNI and hostname verification continue to use the hostname instead of - * an IP literal. Redirects are never followed automatically. The response - * body is capped at {@link #maxResponseBytes()}. + * an IP literal. Redirects are never followed automatically. + * + *

    This is the probe path: the body is capped at the probe cap + * {@link #maxResponseBytes()} (4 KiB by default) and a longer body stops + * being read, with {@link AdcpHttpResponse#truncated()} set. Callers that + * parse the body must reject truncated responses — or use + * {@link #getForAgentCard}, which enforces the strict call cap instead. * * @param method HTTP method (GET, POST, etc.) * @param uri target URI @@ -140,6 +193,10 @@ public AdcpHttpResponse send( requestBuilder.method(method, HttpRequest.BodyPublishers.noBody()); } + // Always ask for an uncompressed body so the cap counts the bytes we + // actually buffer. The JDK HttpClient never decompresses transparently. + requestBuilder.setHeader("Accept-Encoding", "identity"); + // Step 3: Send with body-capped response handler HttpResponse response = httpClient.send( requestBuilder.build(), @@ -148,7 +205,7 @@ public AdcpHttpResponse send( // Step 4: Read body with cap enforcement. // Ensure the InputStream is closed even if readBodyWithCap throws. try { - return readBodyWithCap(response); + return readBodyWithCap(response, maxResponseBytes, false, validatedUri); } catch (Throwable t) { try { response.body().close(); @@ -177,13 +234,40 @@ public AdcpHttpResponse get(URI uri, Map headers) * forwarding risky (transparent redirect-follow, DNS rebinding) are neutralised * by the client's {@code followRedirects(NEVER)} policy and SSRF validation. * + *

    The body is parsed by the caller, so it is held to the strict call cap + * ({@link #maxCallResponseBytes()}): an over-cap card throws + * {@link org.adcontextprotocol.adcp.error.ResponseTooLargeError} rather than + * arriving truncated. + * * @param uri target URI (SSRF-validated) * @param headers headers to forward, including any auth material - * @return the response, body capped at {@link #maxResponseBytes()} + * @return the response, body capped at {@link #maxCallResponseBytes()} + * @throws org.adcontextprotocol.adcp.error.ResponseTooLargeError if the card + * exceeds {@link #maxCallResponseBytes()} */ public AdcpHttpResponse getForAgentCard(URI uri, Map headers) throws IOException, InterruptedException { + return getForAgentCard(uri, headers, maxCallResponseBytes); + } + + /** + * Fetches an agent card with an explicit response cap, so a per-agent or + * per-call override applies to discovery as well as to the call itself. + * + * @param uri target URI (SSRF-validated) + * @param headers headers to forward, including any auth material + * @param maxResponseBytes response cap in bytes; must be positive + * @throws org.adcontextprotocol.adcp.error.ResponseTooLargeError if the card + * exceeds {@code maxResponseBytes} + */ + public AdcpHttpResponse getForAgentCard(URI uri, Map headers, + long maxResponseBytes) + throws IOException, InterruptedException { Objects.requireNonNull(uri, "uri"); + if (maxResponseBytes <= 0) { + throw new IllegalArgumentException( + "maxResponseBytes must be positive: " + maxResponseBytes); + } if (requireHttps && "http".equalsIgnoreCase(uri.getScheme())) { String host = uri.getHost(); @@ -212,6 +296,7 @@ public AdcpHttpResponse getForAgentCard(URI uri, Map headers) }); } + requestBuilder.setHeader("Accept-Encoding", "identity"); requestBuilder.method("GET", HttpRequest.BodyPublishers.noBody()); HttpResponse response = httpClient.send( @@ -219,7 +304,7 @@ public AdcpHttpResponse getForAgentCard(URI uri, Map headers) HttpResponse.BodyHandlers.ofInputStream()); try { - return readBodyWithCap(response); + return readBodyWithCap(response, maxResponseBytes, true, validatedUri); } catch (Throwable t) { try { response.body().close(); @@ -243,27 +328,72 @@ public SsrfPolicy ssrfPolicy() { return ssrfPolicy; } - /** Maximum response body size in bytes. */ + /** Maximum probe response body size in bytes (see {@link #send}). */ public long maxResponseBytes() { return maxResponseBytes; } /** - * Creates an MCP transport client builder with the same connection-timeout - * and redirect policy used by this client. + * Maximum agent-response body size in bytes for paths whose body is parsed: + * agent-card fetches and the MCP/A2A call transports. Default 4 MiB. + */ + public long maxCallResponseBytes() { + return maxCallResponseBytes; + } + + /** + * Creates an MCP transport client builder with the same connection-timeout, + * redirect policy and response cap used by this client. */ public HttpClient.Builder newMcpClientBuilder() { return newHttpClientBuilder(); } /** - * Creates an HTTP client builder with this client's connection-timeout and - * redirect policy ({@code NEVER}). Suitable for any transport (MCP, A2A, etc.). + * Creates an MCP transport client builder with an explicit response cap. + * + * @param maxResponseBytes per-connection response cap in bytes + */ + public HttpClient.Builder newMcpClientBuilder(long maxResponseBytes) { + return newHttpClientBuilder(maxResponseBytes); + } + + /** + * Creates an HTTP client builder with this client's connection-timeout, + * redirect policy ({@code NEVER}) and call-response cap + * ({@link #maxCallResponseBytes()}). Suitable for any transport + * (MCP, A2A, etc.). */ public HttpClient.Builder newHttpClientBuilder() { - return HttpClient.newBuilder() - .connectTimeout(connectTimeout) - .followRedirects(HttpClient.Redirect.NEVER); + return newHttpClientBuilder(maxCallResponseBytes); + } + + /** + * Creates an HTTP client builder with an explicit response cap, for + * transports that need a per-agent or per-call override. + * + *

    The returned builder produces a + * {@link ResponseSizeLimitingHttpClient}: every response it reads is + * counted as it streams and aborted with + * {@link org.adcontextprotocol.adcp.error.ResponseTooLargeError} once the + * cap is crossed. One shape is treated specially — a {@code GET} answered + * with {@code text/event-stream} is a long-lived listening stream and gets + * a larger cumulative budget with the same per-event cap. See + * {@link ResponseSizeLimitingHttpClient} for why, and for what that means + * if you hand this builder to a transport of your own. + * + * @param maxResponseBytes response cap in bytes; must be positive + */ + public HttpClient.Builder newHttpClientBuilder(long maxResponseBytes) { + if (maxResponseBytes <= 0) { + throw new IllegalArgumentException( + "maxResponseBytes must be positive: " + maxResponseBytes); + } + return new ResponseSizeLimitingHttpClient.LimitingClientBuilder( + HttpClient.newBuilder() + .connectTimeout(connectTimeout) + .followRedirects(HttpClient.Redirect.NEVER), + maxResponseBytes); } @Override @@ -333,41 +463,114 @@ private static boolean isLoopback(String host) { || "::1".equals(host); } - private AdcpHttpResponse readBodyWithCap(HttpResponse response) + /** + * Reads the body, counting bytes as they arrive and stopping at {@code cap}. + * + *

    A declared {@code Content-Length} over the cap is rejected before any + * body byte is read. In {@code strict} mode an over-cap body throws + * {@link ResponseTooLargeError}; otherwise (probe paths, where the body is + * not parsed) reading stops at the cap and the result is flagged via + * {@link AdcpHttpResponse#truncated()}. + * + *

    The read is also bounded in time. {@code HttpRequest.timeout} + * only covers the status line and headers, so without this a peer that + * sends headers and then drips — or stalls outright — would pin the calling + * thread and its socket for as long as it liked. A watchdog closes the + * stream at {@code readTimeout}, which unblocks a stalled {@code read()}, + * and the loop checks the same deadline so a drip-feed gets a clear error. + */ + private AdcpHttpResponse readBodyWithCap(HttpResponse response, + long cap, boolean strict, URI uri) throws IOException { - long cap = maxResponseBytes; + if (strict) { + long declared = declaredContentLength(response); + if (declared > cap) { + throw new ResponseTooLargeError(cap, declared, uri); + } + } + boolean truncated = false; long totalRead = 0; + long deadline = System.nanoTime() + readTimeout.toNanos(); try (InputStream is = response.body()) { - ByteArrayOutputStream baos = new ByteArrayOutputStream( - (int) Math.min(cap, 8192)); - byte[] buf = new byte[8192]; - int n; - while ((n = is.read(buf)) != -1) { - totalRead += n; - if (totalRead <= cap) { - baos.write(buf, 0, n); - } else if (!truncated) { - // Write only the portion up to the cap + // Runs on the JDK's shared delayed-task thread; cancelling before it + // fires stops it from running at all. + CompletableFuture watchdog = CompletableFuture.runAsync( + () -> closeQuietly(is), + CompletableFuture.delayedExecutor( + readTimeout.toMillis(), TimeUnit.MILLISECONDS)); + try { + ByteArrayOutputStream baos = new ByteArrayOutputStream( + (int) Math.min(cap, 8192)); + byte[] buf = new byte[8192]; + int n; + while ((n = is.read(buf)) != -1) { + if (System.nanoTime() - deadline > 0) { + throw new HttpTimeoutException( + "Response body from " + redact(uri) + " not delivered within " + + readTimeout + " (read " + totalRead + " bytes)"); + } + totalRead += n; + if (totalRead <= cap) { + baos.write(buf, 0, n); + continue; + } + if (strict) { + throw new ResponseTooLargeError(cap, totalRead, uri); + } + // Probe path: keep the bytes up to the cap, flag, stop reading. int remaining = (int) (cap - (totalRead - n)); if (remaining > 0) { baos.write(buf, 0, remaining); } truncated = true; - log.debug("Response body truncated at {} bytes (cap={})", + log.debug("Probe response body truncated at {} bytes (cap={})", totalRead, cap); - // Stop reading — don't consume the rest break; } + + return new AdcpHttpResponse( + response.statusCode(), + response.headers(), + baos.toByteArray(), + truncated, + totalRead); + } finally { + watchdog.cancel(false); } + } + } + + private static void closeQuietly(InputStream stream) { + try { + stream.close(); + } catch (IOException e) { + log.debug("Error closing response body on read timeout: {}", e.toString()); + } + } + + /** + * Strips userinfo and query from a URI before it lands in an exception + * message or a log line — both routinely carry credentials. + */ + static URI redact(URI uri) { + if (uri.getRawUserInfo() == null && uri.getRawQuery() == null) { + return uri; + } + try { + return new URI(uri.getScheme(), null, uri.getHost(), uri.getPort(), + uri.getPath(), null, null); + } catch (URISyntaxException e) { + return URI.create(String.valueOf(uri.getScheme()) + "://" + uri.getHost()); + } + } - return new AdcpHttpResponse( - response.statusCode(), - response.headers(), - baos.toByteArray(), - truncated, - totalRead); + private static long declaredContentLength(HttpResponse response) { + try { + return response.headers().firstValueAsLong("content-length").orElse(-1L); + } catch (NumberFormatException e) { + return -1L; } } @@ -380,6 +583,7 @@ private static boolean isProtectedHeader(String name) { public static final class Builder { private SsrfPolicy ssrfPolicy = SsrfPolicy.strict(); private long maxResponseBytes = DEFAULT_MAX_RESPONSE_BYTES; + private long maxCallResponseBytes = DEFAULT_MAX_CALL_RESPONSE_BYTES; private Duration connectTimeout = DEFAULT_CONNECT_TIMEOUT; private Duration readTimeout = DEFAULT_READ_TIMEOUT; private String userAgent = DEFAULT_USER_AGENT; @@ -398,20 +602,45 @@ public Builder ssrfPolicy(SsrfPolicy ssrfPolicy) { } /** - * Maximum response body size in bytes. Responses exceeding this - * are truncated and flagged via {@link AdcpHttpResponse#truncated()}. + * Maximum probe response body size in bytes. Probe responses + * exceeding this stop being read and are flagged via + * {@link AdcpHttpResponse#truncated()}. * Default: {@value #DEFAULT_MAX_RESPONSE_BYTES} (4 KiB). - * Maximum: 64 MB. + * Maximum: 64 MiB. + * + *

    This is not the cap on agent responses — that is + * {@link #maxCallResponseBytes(long)}. */ public Builder maxResponseBytes(long maxResponseBytes) { - if (maxResponseBytes <= 0 || maxResponseBytes > 64 * 1024 * 1024) { + if (maxResponseBytes <= 0 || maxResponseBytes > MAX_PROBE_RESPONSE_BYTES) { throw new IllegalArgumentException( - "maxResponseBytes must be in (0, 67108864]: " + maxResponseBytes); + "maxResponseBytes must be in (0, " + MAX_PROBE_RESPONSE_BYTES + + "]: " + maxResponseBytes); } this.maxResponseBytes = maxResponseBytes; return this; } + /** + * Maximum agent-response body size in bytes for paths whose body is + * parsed: agent-card fetches and the MCP/A2A call transports. + * Over-cap responses throw + * {@link org.adcontextprotocol.adcp.error.ResponseTooLargeError}. + * + *

    Default: 4 MiB ({@link #DEFAULT_MAX_CALL_RESPONSE_BYTES}). Raise it + * only for agents that legitimately publish very large catalogues — + * agent responses are untrusted input and the cap is what keeps one + * peer from exhausting the caller's memory. + */ + public Builder maxCallResponseBytes(long maxCallResponseBytes) { + if (maxCallResponseBytes <= 0) { + throw new IllegalArgumentException( + "maxCallResponseBytes must be positive: " + maxCallResponseBytes); + } + this.maxCallResponseBytes = maxCallResponseBytes; + return this; + } + /** Connection timeout. Default: 10 seconds. */ public Builder connectTimeout(Duration connectTimeout) { this.connectTimeout = Objects.requireNonNull(connectTimeout); diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/http/ResponseSizeLimitingHttpClient.java b/adcp/src/main/java/org/adcontextprotocol/adcp/http/ResponseSizeLimitingHttpClient.java new file mode 100644 index 0000000..86637af --- /dev/null +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/http/ResponseSizeLimitingHttpClient.java @@ -0,0 +1,653 @@ +package org.adcontextprotocol.adcp.http; + +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.net.Authenticator; +import java.net.CookieHandler; +import java.net.InetAddress; +import java.net.ProxySelector; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.net.http.WebSocket; +import java.nio.ByteBuffer; +import java.time.Duration; +import java.util.List; +import java.util.Locale; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.Executor; +import java.util.concurrent.Flow; +import java.util.concurrent.atomic.AtomicBoolean; + +import javax.net.ssl.SSLContext; +import javax.net.ssl.SSLParameters; + +/** + * A {@link HttpClient} decorator that enforces a response-size cap + * while the body streams in. + * + *

    Agent responses are untrusted input. Without a cap, a single peer can + * exhaust the caller's memory by answering a tool call with a multi-gigabyte + * body. This decorator exists because the MCP + * ({@code HttpClientStreamableHttpTransport} / {@code HttpClientSseClientTransport}) + * and A2A ({@code JdkA2AHttpClient}) transports accept a JDK {@link HttpClient} + * but expose no body-size hook of their own — wrapping the client is the only + * choke point that covers every call they make. + * + *

    Enforcement: + *

      + *
    1. Pre-reject on {@code Content-Length}. A declared length above + * the cap fails before a single body byte is read.
      2. + *
    2. Count while reading. Bytes are counted in + * {@link Flow.Subscriber#onNext}; the subscription is cancelled the + * moment the cap is crossed, so the oversized body is never buffered.
      3. + *
    3. SSE per event and per call. For + * {@code Content-Type: text/event-stream}, bytes are counted per event + * (reset at each blank-line delimiter, accepting CR, LF and CRLF framing) + * so one huge event cannot escape the cap, and cumulatively for + * the whole response. Answering a call with {@code text/event-stream} + * instead of {@code application/json} therefore buys no extra budget, and + * neither does drip-feeding small events.
      4. + *
    4. Decompressed bytes. Outbound requests are rewritten to + * {@code Accept-Encoding: identity} so the counted wire bytes are the + * bytes the caller will hold. The JDK {@link HttpClient} never + * transparently decompresses, so a peer cannot hide a 1 GiB payload + * behind a 1 MiB gzip stream.
      5. + *
    + * + *

    Long-lived listening streams

    + * + *

    One response shape gets a larger cumulative budget: a {@code GET} answered + * with {@code text/event-stream}. That is the MCP SSE session channel, which + * stays open for the life of the connection and carries every server-to-client + * message, so holding its running total to a single call's cap would kill + * healthy long sessions. Its per-event cap is unchanged — and per-event is the + * cap that bounds memory there, because the transports aggregate one event at a + * time and discard the bytes in between. Its cumulative budget is + * {@code maxResponseBytes * }{@value #LISTENING_STREAM_MULTIPLIER}, finite so + * that a peer cannot hold a {@code GET} open indefinitely simply by choosing an + * SSE content type. + * + *

    Reporting

    + * + *

    Oversized responses fail the response future with + * {@link ResponseTooLargeError} — never truncated, never silently dropped. + * Because the JDK may wrap a body-subscriber failure before it reaches the + * caller, use {@link ResponseTooLargeError#findIn(Throwable)} to recover the + * typed error from a cause chain. + * + *

    With a streaming body handler ({@code BodyHandlers.ofInputStream}, + * {@code ofPublisher}) the response future completes normally at the headers and + * the error surfaces on the first read past the cap instead. The SDK's own paths + * use aggregating handlers; callers wrapping their own client should keep that + * in mind. + * + *

    Scope. The cap is per in-flight response. It bounds what one peer + * can make this process hold for one call — it is not a global memory budget, + * so a caller fanning out across many agents still has to bound its own + * concurrency. + * + * @see AdcpHttpClient#newHttpClientBuilder(long) + * @see ResponseTooLargeError + */ +public final class ResponseSizeLimitingHttpClient extends HttpClient { + + private static final Logger log = + LoggerFactory.getLogger(ResponseSizeLimitingHttpClient.class); + + /** + * How much more cumulative traffic a long-lived {@code GET} event stream may + * carry than a single call response, before it is torn down anyway. + * + *

    The MCP SSE session channel stays open for the life of the connection, + * so capping its running total at the per-call cap would kill healthy long + * sessions. It still gets a finite ceiling: without one, a peer could hold + * a {@code GET} open forever by answering with {@code text/event-stream}, + * and that choice would be the peer's to make, not ours. + */ + public static final int LISTENING_STREAM_MULTIPLIER = 256; + + private static final String SSE_CONTENT_TYPE = "text/event-stream"; + private static final String ACCEPT_ENCODING = "Accept-Encoding"; + private static final String IDENTITY = "identity"; + + private final HttpClient delegate; + private final long maxResponseBytes; + + private ResponseSizeLimitingHttpClient(HttpClient delegate, long maxResponseBytes) { + this.delegate = Objects.requireNonNull(delegate, "delegate"); + if (maxResponseBytes <= 0) { + throw new IllegalArgumentException( + "maxResponseBytes must be positive: " + maxResponseBytes); + } + this.maxResponseBytes = maxResponseBytes; + } + + /** + * Wraps {@code delegate} so that every response it produces is capped at + * {@code maxResponseBytes}. + * + *

    The cap applies per response: to the whole body, to each SSE event, + * and cumulatively across a response's events. A {@code GET} answered with + * {@code text/event-stream} is treated as a long-lived listening stream and + * gets {@value #LISTENING_STREAM_MULTIPLIER}× the cumulative budget, with + * its per-event cap unchanged — see the class javadoc. + * + * @param delegate the client to decorate + * @param maxResponseBytes cap in bytes; must be positive + * @return the decorated client + */ + public static ResponseSizeLimitingHttpClient wrap(HttpClient delegate, + long maxResponseBytes) { + return new ResponseSizeLimitingHttpClient(delegate, maxResponseBytes); + } + + /** The cap in bytes applied to every response from this client. */ + public long maxResponseBytes() { + return maxResponseBytes; + } + + // -- HttpClient: the three methods that actually carry a body -- + + @Override + public HttpResponse send(HttpRequest request, + HttpResponse.BodyHandler responseBodyHandler) + throws IOException, InterruptedException { + return delegate.send(identityEncoding(request), + limiting(request, responseBodyHandler)); + } + + @Override + public CompletableFuture> sendAsync( + HttpRequest request, HttpResponse.BodyHandler responseBodyHandler) { + return delegate.sendAsync(identityEncoding(request), + limiting(request, responseBodyHandler)); + } + + @Override + public CompletableFuture> sendAsync( + HttpRequest request, + HttpResponse.BodyHandler responseBodyHandler, + HttpResponse.@Nullable PushPromiseHandler pushPromiseHandler) { + return delegate.sendAsync(identityEncoding(request), + limiting(request, responseBodyHandler), + limiting(pushPromiseHandler)); + } + + // -- HttpClient: pure delegation -- + + @Override + public Optional cookieHandler() { + return delegate.cookieHandler(); + } + + @Override + public Optional connectTimeout() { + return delegate.connectTimeout(); + } + + @Override + public Redirect followRedirects() { + return delegate.followRedirects(); + } + + @Override + public Optional proxy() { + return delegate.proxy(); + } + + @Override + public SSLContext sslContext() { + return delegate.sslContext(); + } + + @Override + public SSLParameters sslParameters() { + return delegate.sslParameters(); + } + + @Override + public Optional authenticator() { + return delegate.authenticator(); + } + + @Override + public Version version() { + return delegate.version(); + } + + @Override + public Optional executor() { + return delegate.executor(); + } + + @Override + public WebSocket.Builder newWebSocketBuilder() { + return delegate.newWebSocketBuilder(); + } + + @Override + public void shutdown() { + delegate.shutdown(); + } + + @Override + public boolean awaitTermination(Duration duration) throws InterruptedException { + return delegate.awaitTermination(duration); + } + + @Override + public boolean isTerminated() { + return delegate.isTerminated(); + } + + @Override + public void shutdownNow() { + delegate.shutdownNow(); + } + + @Override + public void close() { + delegate.close(); + } + + @Override + public String toString() { + return "ResponseSizeLimitingHttpClient[maxResponseBytes=" + maxResponseBytes + + ", delegate=" + delegate + "]"; + } + + // -- internal -- + + /** + * Rewrites the request to ask for an uncompressed body, so the bytes this + * client counts are the bytes the caller ends up holding. + */ + private static HttpRequest identityEncoding(HttpRequest request) { + String current = request.headers().firstValue(ACCEPT_ENCODING).orElse(null); + if (IDENTITY.equalsIgnoreCase(current)) { + return request; + } + if (current != null) { + log.debug("Replacing Accept-Encoding '{}' with 'identity' so the response" + + " size cap counts uncompressed bytes", current); + } + return HttpRequest.newBuilder(request, + (name, value) -> !ACCEPT_ENCODING.equalsIgnoreCase(name)) + .header(ACCEPT_ENCODING, IDENTITY) + .build(); + } + + private HttpResponse.BodyHandler limiting( + HttpRequest request, HttpResponse.BodyHandler delegateHandler) { + Objects.requireNonNull(delegateHandler, "responseBodyHandler"); + // A GET that comes back as SSE is a long-lived listening stream, not the + // response to a single call, so its running total gets a larger ceiling. + boolean listeningStream = "GET".equalsIgnoreCase(request.method()); + URI uri = request.uri(); + long cap = maxResponseBytes; + return responseInfo -> { + // A declared Content-Length is held to the per-call cap even on a + // GET event stream: a body of known, finite size is a payload, not a + // listening stream, so it gets no share of the larger ceiling. + long declared = declaredContentLength(responseInfo); + if (declared > cap) { + log.debug("Rejecting response from {}: Content-Length {} exceeds cap {}", + redact(uri), declared, cap); + return new FailedBodySubscriber<>( + new ResponseTooLargeError(cap, declared, uri)); + } + boolean sse = isEventStream(responseInfo); + long totalCap = sse && listeningStream + ? saturatingMultiply(cap, LISTENING_STREAM_MULTIPLIER) + : cap; + return new LimitingBodySubscriber<>( + delegateHandler.apply(responseInfo), cap, totalCap, sse, uri); + }; + } + + private HttpResponse.@Nullable PushPromiseHandler limiting( + HttpResponse.@Nullable PushPromiseHandler delegateHandler) { + if (delegateHandler == null) { + return null; + } + return (initiatingRequest, pushPromiseRequest, acceptor) -> + delegateHandler.applyPushPromise(initiatingRequest, pushPromiseRequest, + handler -> acceptor.apply(limiting(pushPromiseRequest, handler))); + } + + /** Keeps credentials in a URI out of log lines. */ + private static Object redact(@Nullable URI uri) { + return uri == null ? "(unknown)" : AdcpHttpClient.redact(uri); + } + + private static long saturatingMultiply(long value, int factor) { + try { + return Math.multiplyExact(value, (long) factor); + } catch (ArithmeticException overflow) { + return Long.MAX_VALUE; + } + } + + private static long declaredContentLength(HttpResponse.ResponseInfo responseInfo) { + try { + return responseInfo.headers().firstValueAsLong("content-length").orElse(-1L); + } catch (NumberFormatException e) { + // Malformed Content-Length — treat as unknown and let the streaming + // counter do the work. + return -1L; + } + } + + private static boolean isEventStream(HttpResponse.ResponseInfo responseInfo) { + return responseInfo.headers().firstValue("content-type") + .map(ct -> ct.toLowerCase(Locale.ROOT).contains(SSE_CONTENT_TYPE)) + .orElse(false); + } + + /** + * Fails the response immediately without reading any body — used when + * {@code Content-Length} already declares an oversized payload. + */ + private static final class FailedBodySubscriber + implements HttpResponse.BodySubscriber { + + private final Throwable error; + + FailedBodySubscriber(Throwable error) { + this.error = error; + } + + @Override + public void onSubscribe(Flow.Subscription subscription) { + subscription.cancel(); + } + + @Override + public void onNext(List item) { + // never requested + } + + @Override + public void onError(Throwable throwable) { + // already failed + } + + @Override + public void onComplete() { + // already failed + } + + @Override + public CompletionStage getBody() { + return CompletableFuture.failedFuture(error); + } + } + + /** + * Counts bytes as the body streams in and aborts the subscription the + * moment the cap is crossed. + * + *

    The reactive-streams contract guarantees {@code onNext} is signalled + * serially, so the counters need no locking — but successive signals may + * arrive on different carrier threads, so they are {@code volatile} for + * visibility. The per-byte scan works on locals and writes back once per + * {@code onNext} to keep that cheap. The terminal-state flag is atomic + * because {@code onError} may race with a cap breach. + */ + private static final class LimitingBodySubscriber + implements HttpResponse.BodySubscriber { + + private final HttpResponse.BodySubscriber delegate; + private final long maxBytes; + private final long maxTotalBytes; + private final boolean sse; + private final @Nullable URI uri; + private final CompletableFuture body = new CompletableFuture<>(); + private final AtomicBoolean aborted = new AtomicBoolean(); + + private volatile Flow.@Nullable Subscription subscription; + + private volatile long totalBytes; + private volatile long eventBytes; + private volatile int terminators; + private volatile boolean pendingCr; + + LimitingBodySubscriber(HttpResponse.BodySubscriber delegate, long maxBytes, + long maxTotalBytes, boolean sse, @Nullable URI uri) { + this.delegate = Objects.requireNonNull(delegate, "delegate"); + this.maxBytes = maxBytes; + this.maxTotalBytes = maxTotalBytes; + this.sse = sse; + this.uri = uri; + delegate.getBody().whenComplete((value, failure) -> { + if (failure != null) { + body.completeExceptionally(failure); + } else { + body.complete(value); + } + }); + } + + @Override + public void onSubscribe(Flow.Subscription subscription) { + this.subscription = subscription; + delegate.onSubscribe(subscription); + } + + @Override + public void onNext(List buffers) { + if (aborted.get()) { + return; + } + for (ByteBuffer buffer : buffers) { + totalBytes += buffer.remaining(); + if (totalBytes > maxTotalBytes) { + abort(maxTotalBytes, totalBytes); + return; + } + if (sse && scanForOversizedEvent(buffer)) { + abort(maxBytes, eventBytes); + return; + } + } + delegate.onNext(buffers); + } + + @Override + public void onError(Throwable throwable) { + if (aborted.get()) { + return; + } + delegate.onError(throwable); + } + + @Override + public void onComplete() { + if (aborted.get()) { + return; + } + delegate.onComplete(); + } + + @Override + public CompletionStage getBody() { + return body; + } + + /** + * Counts bytes since the last SSE event delimiter without consuming the + * buffer (absolute gets only — the delegate still sees the full body). + * + * @return {@code true} if the current event has crossed the cap + */ + private boolean scanForOversizedEvent(ByteBuffer buffer) { + long bytes = eventBytes; + int runOfTerminators = terminators; + boolean cr = pendingCr; + boolean oversized = false; + + for (int i = buffer.position(); i < buffer.limit(); i++) { + byte b = buffer.get(i); + bytes++; + + // SSE accepts CR, LF and CRLF as line terminators. A pending CR + // always terminated its line; whether the current byte joins it + // (LF), starts another terminator (CR), or is content is a + // separate question — hence two independent flags. + boolean endedLine; + boolean contentByte; + if (cr) { + cr = b == '\r'; + endedLine = true; // the pending CR ended a line + contentByte = b != '\n' && b != '\r'; + } else if (b == '\r') { + cr = true; + endedLine = false; // wait for the next byte + contentByte = false; + } else { + endedLine = b == '\n'; + contentByte = !endedLine; + } + + boolean eventEnded = false; + if (endedLine && ++runOfTerminators >= 2) { + // Two terminators in a row is a blank line: end of event. + runOfTerminators = 0; + bytes = contentByte ? 1 : 0; // the content byte opens the next event + eventEnded = true; + } + if (contentByte) { + runOfTerminators = 0; + } + + if (!eventEnded && bytes > maxBytes) { + oversized = true; + break; + } + } + + eventBytes = bytes; + terminators = runOfTerminators; + pendingCr = cr; + return oversized; + } + + private void abort(long limit, long bytesRead) { + if (!aborted.compareAndSet(false, true)) { + return; + } + ResponseTooLargeError error = new ResponseTooLargeError(limit, bytesRead, uri); + log.debug("Aborting response from {}: {}", redact(uri), error.getMessage()); + // Complete our own future first so the delegate's terminal signal + // cannot overwrite the typed error. + body.completeExceptionally(error); + Flow.Subscription current = subscription; + if (current != null) { + current.cancel(); + } + try { + delegate.onError(error); + } catch (RuntimeException e) { + log.debug("Delegate body subscriber rejected the abort signal: {}", + e.toString()); + } + } + } + + /** + * An {@link HttpClient.Builder} that produces size-limited clients. + * + *

    Lets a size cap be handed to third-party transports that only accept a + * {@code HttpClient.Builder} (the MCP transports do exactly this). + */ + static final class LimitingClientBuilder implements HttpClient.Builder { + + private final HttpClient.Builder delegate; + private final long maxResponseBytes; + + LimitingClientBuilder(HttpClient.Builder delegate, long maxResponseBytes) { + this.delegate = Objects.requireNonNull(delegate, "delegate"); + this.maxResponseBytes = maxResponseBytes; + } + + @Override + public HttpClient.Builder cookieHandler(CookieHandler cookieHandler) { + delegate.cookieHandler(cookieHandler); + return this; + } + + @Override + public HttpClient.Builder connectTimeout(Duration duration) { + delegate.connectTimeout(duration); + return this; + } + + @Override + public HttpClient.Builder sslContext(SSLContext sslContext) { + delegate.sslContext(sslContext); + return this; + } + + @Override + public HttpClient.Builder sslParameters(SSLParameters sslParameters) { + delegate.sslParameters(sslParameters); + return this; + } + + @Override + public HttpClient.Builder executor(Executor executor) { + delegate.executor(executor); + return this; + } + + @Override + public HttpClient.Builder followRedirects(Redirect policy) { + delegate.followRedirects(policy); + return this; + } + + @Override + public HttpClient.Builder version(Version version) { + delegate.version(version); + return this; + } + + @Override + public HttpClient.Builder priority(int priority) { + delegate.priority(priority); + return this; + } + + @Override + public HttpClient.Builder proxy(ProxySelector proxySelector) { + delegate.proxy(proxySelector); + return this; + } + + @Override + public HttpClient.Builder authenticator(Authenticator authenticator) { + delegate.authenticator(authenticator); + return this; + } + + @Override + public HttpClient.Builder localAddress(InetAddress localAddr) { + delegate.localAddress(localAddr); + return this; + } + + @Override + public HttpClient build() { + return wrap(delegate.build(), maxResponseBytes); + } + } +} diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/CallToolOptions.java b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/CallToolOptions.java index 0c6de00..8800df7 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/CallToolOptions.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/CallToolOptions.java @@ -7,15 +7,24 @@ /** * Options for a single {@code callTool()} invocation. * - *

    Implementation status: In v0.1, MCP transport applies a - * fixed 10 MB content limit regardless of {@code maxResponseBytes}. - * The {@code timeout} and {@code maxResponseBytes} fields are accepted - * for forward compatibility but are not yet wired into the - * MCP transport path. They will be enforced when the call-level timeout - * and per-agent body-cap features ship (planned v0.2). + *

    {@code maxResponseBytes} overrides the response size cap for this call on + * both the MCP and A2A transports. It is enforced while the body streams in; + * an over-cap response throws + * {@link org.adcontextprotocol.adcp.error.ResponseTooLargeError}. Resolution + * order is call → agent ({@code AgentConfig.maxResponseBytes}) → client + * ({@code AdcpClient.Builder.maxResponseBytes}, default 4 MiB). + * + *

    Overriding the cap selects a transport connection built for that cap, so + * varying it per call fragments the connection cache. Prefer setting it once on + * the {@link org.adcontextprotocol.adcp.AgentConfig} for agents that routinely + * answer with large catalogues. + * + *

    Implementation status: {@code timeout} is accepted for forward + * compatibility but is not yet wired into the transport path; the + * client-level {@code requestTimeout} applies instead. * * @param timeout per-call timeout (overrides client default) — reserved, not yet enforced - * @param maxResponseBytes per-call body cap (overrides client default) — reserved, not yet enforced + * @param maxResponseBytes per-call response cap in bytes (overrides agent and client defaults) * @param validateResponse whether to validate the response against schema */ public record CallToolOptions( @@ -24,7 +33,10 @@ public record CallToolOptions( boolean validateResponse ) { - /** Default options: no timeout override, no body cap override, validation off. */ + /** + * Default options: no timeout override, no cap override (the agent or + * client default applies), validation off. + */ public static final CallToolOptions DEFAULT = new CallToolOptions(null, null, false); public static Builder builder() { @@ -43,6 +55,11 @@ public Builder timeout(Duration timeout) { return this; } + /** + * Caps this call's response at {@code maxResponseBytes}, overriding the + * agent and client defaults. Must be positive — the cap cannot be + * disabled. + */ public Builder maxResponseBytes(long maxResponseBytes) { if (maxResponseBytes <= 0) { throw new IllegalArgumentException( diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/ProtocolClient.java b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/ProtocolClient.java index 5a3bbe6..a29e620 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/ProtocolClient.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/ProtocolClient.java @@ -7,6 +7,7 @@ import org.adcontextprotocol.adcp.Protocol; import org.adcontextprotocol.adcp.auth.AuthTokenResolver; import org.adcontextprotocol.adcp.error.ProtocolError; +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; import org.adcontextprotocol.adcp.http.ProtectedHeaders; import org.adcontextprotocol.adcp.http.SsrfPolicy; import org.adcontextprotocol.adcp.transport.a2a.A2aCaller; @@ -86,9 +87,12 @@ public T callTool(AgentConfig agent, String toolName, // 1. Validate agent URL against SSRF policy validateUrl(agent); - // 2. Warn if non-default options are passed (not yet enforced in v0.1) - if (!options.equals(CallToolOptions.DEFAULT)) { - log.debug("CallToolOptions fields are not yet enforced by the MCP transport (v0.1)"); + // 2. Resolve the response size cap: call override, then agent, then + // the client default baked into the connection managers. + Long maxResponseBytes = resolveMaxResponseBytes(agent, options); + if (options.timeout() != null) { + log.debug("CallToolOptions.timeout is not yet enforced by the transports; " + + "the client-level requestTimeout applies"); } // 3. Resolve auth headers @@ -113,8 +117,22 @@ public T callTool(AgentConfig agent, String toolName, // 6. Dispatch to transport return agent.protocol() == Protocol.A2A - ? callViaA2a(agent, toolName, mergedArgs, allHeaders, responseType) - : callViaMcp(agent, toolName, mergedArgs, allHeaders, responseType); + ? callViaA2a(agent, toolName, mergedArgs, allHeaders, responseType, + maxResponseBytes) + : callViaMcp(agent, toolName, mergedArgs, allHeaders, responseType, + maxResponseBytes); + } + + /** + * Resolves the response size cap for a call: explicit call option, then the + * agent's own setting, then {@code null} meaning "use the client default". + */ + private static @Nullable Long resolveMaxResponseBytes(AgentConfig agent, + CallToolOptions options) { + if (options.maxResponseBytes() != null) { + return options.maxResponseBytes(); + } + return agent.maxResponseBytes(); } /** @@ -137,13 +155,17 @@ public void close() { private T callViaMcp(AgentConfig agent, String toolName, Map mergedArgs, Map headers, - Class responseType) { + Class responseType, + @Nullable Long maxResponseBytes) { String cacheHash = computeCacheHash(agent); McpSyncClient client = mcpConnectionManager.getOrConnect( - agent.agentUri(), headers, cacheHash); + agent.agentUri(), headers, cacheHash, maxResponseBytes); + long cap = maxResponseBytes != null + ? maxResponseBytes + : mcpConnectionManager.defaultMaxResponseBytes(); try { - return mcpCaller.callTool(client, toolName, mergedArgs, responseType); + return mcpCaller.callTool(client, toolName, mergedArgs, responseType, cap); } catch (ProtocolError e) { if (!isTransportError(e)) { throw e; @@ -155,9 +177,9 @@ private T callViaMcp(AgentConfig agent, String toolName, ProtocolError original = e; client = mcpConnectionManager.getOrConnect( - agent.agentUri(), headers, cacheHash); + agent.agentUri(), headers, cacheHash, maxResponseBytes); try { - return mcpCaller.callTool(client, toolName, mergedArgs, responseType); + return mcpCaller.callTool(client, toolName, mergedArgs, responseType, cap); } catch (ProtocolError retry) { retry.addSuppressed(original); throw retry; @@ -166,6 +188,11 @@ private T callViaMcp(AgentConfig agent, String toolName, } private boolean isTransportError(ProtocolError e) { + // An oversized response surfaces as an IOException from the JDK HTTP + // stack, but retrying it just re-downloads the same oversized payload. + if (ResponseTooLargeError.findIn(e) != null) { + return false; + } // Walk the full cause chain — any I/O or timeout failure is transient for (Throwable t = e.getCause(); t != null; t = t.getCause()) { if (t instanceof java.io.IOException @@ -179,11 +206,16 @@ private boolean isTransportError(ProtocolError e) { private T callViaA2a(AgentConfig agent, String toolName, Map mergedArgs, Map headers, - Class responseType) { + Class responseType, + @Nullable Long maxResponseBytes) { String cacheHash = computeCacheHash(agent); - var client = a2aConnectionManager.getOrConnect(agent, headers, cacheHash); + var client = a2aConnectionManager.getOrConnect( + agent, headers, cacheHash, maxResponseBytes); + long cap = maxResponseBytes != null + ? maxResponseBytes + : a2aConnectionManager.defaultMaxResponseBytes(); try { - return a2aCaller.callTool(client, toolName, mergedArgs, responseType, headers); + return a2aCaller.callTool(client, toolName, mergedArgs, responseType, headers, cap); } catch (ProtocolError e) { if (!isTransportError(e)) { throw e; @@ -193,9 +225,11 @@ private T callViaA2a(AgentConfig agent, String toolName, toolName, e.getMessage()); ProtocolError original = e; - client = a2aConnectionManager.getOrConnect(agent, headers, cacheHash); + client = a2aConnectionManager.getOrConnect( + agent, headers, cacheHash, maxResponseBytes); try { - return a2aCaller.callTool(client, toolName, mergedArgs, responseType, headers); + return a2aCaller.callTool(client, toolName, mergedArgs, responseType, + headers, cap); } catch (ProtocolError retry) { retry.addSuppressed(original); throw retry; diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/a2a/A2aCaller.java b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/a2a/A2aCaller.java index b8d73eb..989e398 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/a2a/A2aCaller.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/a2a/A2aCaller.java @@ -17,6 +17,8 @@ import org.a2aproject.sdk.spec.TaskState; import org.a2aproject.sdk.spec.TextPart; import org.adcontextprotocol.adcp.error.ProtocolError; +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; +import org.adcontextprotocol.adcp.http.AdcpHttpClient; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -33,11 +35,17 @@ /** * Calls AdCP tools over A2A JSON-RPC + SSE. + * + *

    The authoritative response size cap lives in the transport's HTTP client + * (see {@link org.adcontextprotocol.adcp.http.ResponseSizeLimitingHttpClient}), + * which aborts an oversized body while it streams. The part checks here are a + * backstop, and they translate a transport-level + * {@link ResponseTooLargeError} back into the typed error instead of burying it + * in a {@link ProtocolError}. */ public final class A2aCaller { private static final Logger log = LoggerFactory.getLogger(A2aCaller.class); - private static final int MAX_CONTENT_LENGTH = 10 * 1024 * 1024; private static final int MAX_ERROR_LENGTH = 500; private static final long RESPONSE_TIMEOUT_SECONDS = 30; private static final String TOOL_NAME_KEY = "adcp_tool_name"; @@ -53,14 +61,45 @@ public A2aCaller(ObjectMapper objectMapper) { this.objectMapper.deactivateDefaultTyping(); } + /** + * Calls an AdCP tool over A2A with no extra headers. + * + *

    Uses the SDK's default 4 MiB cap as its content backstop, not the cap + * configured on the calling client — the transport's own cap is the + * authoritative one. Prefer + * {@link #callTool(Client, String, Map, Class, Map, long)}. + */ public T callTool(Client client, String toolName, Map args, Class responseType) { - return callTool(new ClientAdapter(client), toolName, args, responseType, Map.of()); + return callTool(client, toolName, args, responseType, Map.of(), + AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES); } + /** + * Calls an AdCP tool over A2A with extra headers. + * + *

    Uses the SDK's default 4 MiB cap as its content backstop, not the cap + * configured on the calling client — the transport's own cap is the + * authoritative one. Prefer + * {@link #callTool(Client, String, Map, Class, Map, long)}. + */ public T callTool(Client client, String toolName, Map args, Class responseType, Map headers) { + return callTool(client, toolName, args, responseType, headers, + AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES); + } + + /** + * Calls an AdCP tool over A2A, rejecting response parts larger than + * {@code maxResponseBytes}. + * + * @param maxResponseBytes the effective response size cap for this call + * @throws ResponseTooLargeError if the agent's response exceeds the cap + */ + public T callTool(Client client, String toolName, + Map args, Class responseType, + Map headers, long maxResponseBytes) { for (var entry : headers.entrySet()) { String k = entry.getKey(); String v = entry.getValue(); @@ -70,12 +109,20 @@ public T callTool(Client client, String toolName, "callTool headers must not contain CR/LF or null: " + k); } } - return callTool(new ClientAdapter(client), toolName, args, responseType, headers); + return callTool(new ClientAdapter(client), toolName, args, responseType, headers, + maxResponseBytes); } T callTool(A2aMessageClient client, String toolName, Map args, Class responseType, Map headers) { + return callTool(client, toolName, args, responseType, headers, + AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES); + } + + T callTool(A2aMessageClient client, String toolName, + Map args, Class responseType, + Map headers, long maxResponseBytes) { // Validate toolName before use — reject rather than silently mutate the outbound request if (toolName == null || toolName.isBlank()) { throw new IllegalArgumentException("toolName must not be null or blank"); @@ -139,7 +186,8 @@ T callTool(A2aMessageClient client, String toolName, if (failure.get() != null) { throw wrapFailure(safeToolName, failure.get()); } - return extractResponse(latestMessage.get(), latestTask.get(), responseType); + return extractResponse(latestMessage.get(), latestTask.get(), responseType, + maxResponseBytes); } catch (ProtocolError e) { throw e; } catch (A2AClientException e) { @@ -163,9 +211,10 @@ private MessageSendParams buildRequest(String toolName, Map args .build(); } - private T extractResponse(Message message, Task task, Class responseType) { + private T extractResponse(Message message, Task task, Class responseType, + long maxResponseBytes) { if (message != null) { - return extractFromParts(message.parts(), responseType); + return extractFromParts(message.parts(), responseType, maxResponseBytes); } if (task != null) { if (task.status() != null && task.status().state() == TaskState.TASK_STATE_FAILED) { @@ -179,14 +228,16 @@ private T extractResponse(Message message, Task task, Class responseType) null); } if (task.status() != null && task.status().message() != null) { - return extractFromParts(task.status().message().parts(), responseType); + return extractFromParts(task.status().message().parts(), responseType, + maxResponseBytes); } if (task.history() != null && !task.history().isEmpty()) { int limit = Math.min(task.history().size(), MAX_HISTORY_SCAN); for (int i = task.history().size() - 1; i >= task.history().size() - limit; i--) { Message historyMessage = task.history().get(i); if (historyMessage != null && historyMessage.parts() != null && !historyMessage.parts().isEmpty()) { - return extractFromParts(historyMessage.parts(), responseType); + return extractFromParts(historyMessage.parts(), responseType, + maxResponseBytes); } } } @@ -194,7 +245,8 @@ private T extractResponse(Message message, Task task, Class responseType) throw new ProtocolError("a2a", "Empty response from A2A sendMessage", null); } - private T extractFromParts(List> parts, Class responseType) { + private T extractFromParts(List> parts, Class responseType, + long maxResponseBytes) { if (parts == null || parts.isEmpty()) { throw new ProtocolError("a2a", "A2A response message had no parts", null); } @@ -206,14 +258,11 @@ private T extractFromParts(List> parts, Class responseType) { if (part instanceof DataPart dataPart) { try { byte[] bytes = objectMapper.writeValueAsBytes(dataPart.data()); - if (bytes.length > MAX_CONTENT_LENGTH) { - throw new ProtocolError("a2a", - "A2A DataPart response exceeds size limit (" - + bytes.length + " > " + MAX_CONTENT_LENGTH + ")", - null); + if (bytes.length > maxResponseBytes) { + throw new ResponseTooLargeError(maxResponseBytes, bytes.length, null); } return objectMapper.readValue(bytes, responseType); - } catch (ProtocolError e) { + } catch (ProtocolError | ResponseTooLargeError e) { throw e; } catch (Exception e) { if (firstParseError == null) { @@ -227,12 +276,14 @@ private T extractFromParts(List> parts, Class responseType) { if (text == null) { continue; } + // UTF-8 is never fewer bytes than chars, so a char count over + // the cap is already over the cap: reject without encoding. + if (text.length() > maxResponseBytes) { + throw new ResponseTooLargeError(maxResponseBytes, text.length(), null); + } int textBytes = text.getBytes(StandardCharsets.UTF_8).length; - if (textBytes > MAX_CONTENT_LENGTH) { - throw new ProtocolError("a2a", - "A2A response content exceeds size limit (" - + textBytes + " > " + MAX_CONTENT_LENGTH + ")", - null); + if (textBytes > maxResponseBytes) { + throw new ResponseTooLargeError(maxResponseBytes, textBytes, null); } try { return objectMapper.readValue(text, responseType); @@ -264,7 +315,14 @@ private boolean isTerminal(Task task) { && task.status().state().isFinal(); } - private ProtocolError wrapFailure(String toolName, Throwable throwable) { + private RuntimeException wrapFailure(String toolName, Throwable throwable) { + // The cap is enforced inside the JDK body subscriber, so the typed error + // arrives wrapped. Surface it rather than hiding it in a ProtocolError, + // which the retry path would treat as transient. + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(throwable); + if (tooLarge != null) { + return tooLarge; + } String message = throwable.getMessage(); return new ProtocolError("a2a", "A2A sendMessage failed for " + toolName + ": " + sanitizeErrorText(message), diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/a2a/A2aConnectionManager.java b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/a2a/A2aConnectionManager.java index 15606f8..bf720a8 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/a2a/A2aConnectionManager.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/a2a/A2aConnectionManager.java @@ -12,9 +12,11 @@ import org.a2aproject.sdk.spec.AgentInterface; import org.adcontextprotocol.adcp.AgentConfig; import org.adcontextprotocol.adcp.error.ProtocolError; +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; import org.adcontextprotocol.adcp.http.AdcpHttpClient; import org.adcontextprotocol.adcp.http.AdcpHttpResponse; import org.adcontextprotocol.adcp.http.ProtectedHeaders; +import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -23,11 +25,14 @@ import java.net.URLEncoder; import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.HashMap; +import java.util.IdentityHashMap; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.TreeMap; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Semaphore; import java.util.concurrent.locks.ReentrantLock; @@ -37,7 +42,9 @@ * *

    Headers are included in the cache key because agent-card discovery is * header-sensitive. The separate cache hash isolates clients by credentials - * without storing raw secrets in the cache key. + * without storing raw secrets in the cache key. The response size cap is in the + * key too: it is baked into the transport's {@link java.net.http.HttpClient}, + * so a call that overrides the cap must not reuse a client built with another. */ public final class A2aConnectionManager implements AutoCloseable { @@ -52,15 +59,24 @@ public final class A2aConnectionManager implements AutoCloseable { private final Semaphore[] connectStripes; private final AgentCardLoader agentCardLoader; private final ClientFactory clientFactory; + private final long defaultMaxResponseBytes; private volatile boolean closed; public A2aConnectionManager(AdcpHttpClient adcpHttpClient, ObjectMapper objectMapper) { - this(new HttpAgentCardLoader(adcpHttpClient, objectMapper), new DefaultClientFactory(adcpHttpClient)); + this(new HttpAgentCardLoader(adcpHttpClient, objectMapper), + new DefaultClientFactory(adcpHttpClient), + adcpHttpClient.maxCallResponseBytes()); } A2aConnectionManager(AgentCardLoader agentCardLoader, ClientFactory clientFactory) { + this(agentCardLoader, clientFactory, AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES); + } + + A2aConnectionManager(AgentCardLoader agentCardLoader, ClientFactory clientFactory, + long defaultMaxResponseBytes) { this.agentCardLoader = Objects.requireNonNull(agentCardLoader, "agentCardLoader"); this.clientFactory = Objects.requireNonNull(clientFactory, "clientFactory"); + this.defaultMaxResponseBytes = defaultMaxResponseBytes; this.connectStripes = new Semaphore[STRIPE_COUNT]; for (int i = 0; i < STRIPE_COUNT; i++) { connectStripes[i] = new Semaphore(1); @@ -68,14 +84,26 @@ public A2aConnectionManager(AdcpHttpClient adcpHttpClient, ObjectMapper objectMa } public Client getOrConnect(AgentConfig agent, Map headers, String cacheHash) { + return getOrConnect(agent, headers, cacheHash, null); + } + + /** + * Gets or creates a cached A2A client with an explicit response size cap. + * + * @param maxResponseBytes per-client response cap, or {@code null} for the + * client default + */ + public Client getOrConnect(AgentConfig agent, Map headers, String cacheHash, + @Nullable Long maxResponseBytes) { if (closed) { throw new IllegalStateException("A2aConnectionManager is closed"); } Objects.requireNonNull(cacheHash, "cacheHash"); Objects.requireNonNull(headers, "headers"); + long cap = effectiveMaxResponseBytes(maxResponseBytes); Map sanitizedAll = sanitizeHeaders(headers); Map sanitizedForKey = filterProtected(sanitizedAll); - String cacheKey = buildCacheKey(agent.agentUri(), sanitizedForKey, cacheHash); + String cacheKey = buildCacheKey(agent.agentUri(), sanitizedForKey, cacheHash, cap); cacheLock.lock(); try { @@ -109,7 +137,7 @@ public Client getOrConnect(AgentConfig agent, Map headers, Strin cacheLock.unlock(); } - Client client = connect(agent, sanitizedAll); + Client client = connect(agent, sanitizedAll, cap); cacheLock.lock(); try { @@ -136,12 +164,28 @@ public void evict(URI agentUri) { public void evict(URI agentUri, String cacheHash) { Objects.requireNonNull(cacheHash, "cacheHash"); - String prefix = agentUri + "#" + cacheHash; + String prefix = agentUri + "#" + cacheHash + "#"; // Evicts all cache entries for the given agent URI and credential hash, - // regardless of which non-secret discovery headers they used. - // This is intentional: on a transport error, all variants for those - // credentials are assumed stale. - evictMatching(key -> key.equals(prefix) || key.startsWith(prefix + "?")); + // regardless of which non-secret discovery headers or response cap they + // used. This is intentional: on a transport error, all variants for + // those credentials are assumed stale. + evictMatching(key -> key.startsWith(prefix)); + } + + /** The response cap applied when a call does not override it. */ + public long defaultMaxResponseBytes() { + return defaultMaxResponseBytes; + } + + private long effectiveMaxResponseBytes(@Nullable Long override) { + if (override == null) { + return defaultMaxResponseBytes; + } + if (override <= 0) { + throw new IllegalArgumentException( + "maxResponseBytes must be positive: " + override); + } + return override; } @Override @@ -154,14 +198,22 @@ public void close() { } finally { cacheLock.unlock(); } + try { + clientFactory.close(); + } catch (RuntimeException e) { + log.debug("Error closing A2A client factory: {}", sanitizeLogText(e.getMessage())); + } } - private Client connect(AgentConfig agent, Map sanitizedHeaders) { + private Client connect(AgentConfig agent, Map sanitizedHeaders, long cap) { try { - AgentCard card = agentCardLoader.load(agent, sanitizedHeaders); - return clientFactory.create(card); + AgentCard card = agentCardLoader.load(agent, sanitizedHeaders, cap); + return clientFactory.create(card, cap); } catch (ProtocolError e) { throw e; + } catch (ResponseTooLargeError e) { + // An oversized agent card is a hard failure, not a connectivity blip. + throw e; } catch (A2AClientException e) { throw new ProtocolError("a2a", "Failed to create A2A client for " + agent.agentUri(), e); } catch (Exception e) { @@ -176,9 +228,17 @@ private Client connect(AgentConfig agent, Map sanitizedHeaders) * crafted {@code =} or {@code &} characters. */ static String buildCacheKey(URI agentUri, Map sanitizedHeaders, String cacheHash) { + return buildCacheKey(agentUri, sanitizedHeaders, cacheHash, + AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES); + } + + static String buildCacheKey(URI agentUri, Map sanitizedHeaders, String cacheHash, + long maxResponseBytes) { StringBuilder sb = new StringBuilder(agentUri.toString()) .append('#') - .append(cacheHash); + .append(cacheHash) + .append('#') + .append(maxResponseBytes); if (sanitizedHeaders.isEmpty()) { return sb.toString(); } @@ -277,12 +337,22 @@ private static String sanitizeForLog(String s) { } private void closeQuietly(Client client) { + if (client == null) { + return; + } try { - if (client != null) { - client.close(); - } + client.close(); } catch (Exception e) { log.debug("Error closing A2A client: {}", sanitizeLogText(e.getMessage())); + } finally { + try { + // Client.close() only closes the A2A transport; the JDK HTTP + // client underneath it belongs to the factory. + clientFactory.release(client); + } catch (RuntimeException e) { + log.debug("Error releasing A2A transport resources: {}", + sanitizeLogText(e.getMessage())); + } } } @@ -295,42 +365,170 @@ private static String sanitizeLogText(String raw) { } interface AgentCardLoader { - AgentCard load(AgentConfig agent, Map headers); + AgentCard load(AgentConfig agent, Map headers, long maxResponseBytes); } interface ClientFactory { - Client create(AgentCard agentCard) throws A2AClientException; + Client create(AgentCard agentCard, long maxResponseBytes) throws A2AClientException; + + /** + * Releases transport resources the factory allocated for {@code client}. + * Called once per client, when the connection cache lets it go. + */ + default void release(Client client) { + } + + /** Releases anything the factory still owns after every client is gone. */ + default void close() { + } } - private static final class DefaultClientFactory implements ClientFactory { + static final class DefaultClientFactory implements ClientFactory { + + private final AdcpHttpClient adcpHttpClient; + /** - * SSRF-safe HTTP client used by the A2A JSON-RPC transport. - * Backed by the same {@link java.net.http.HttpClient} that was built with - * {@code followRedirects(NEVER)} and the configured connect timeout, so - * the transport cannot follow HTTP redirects to internal addresses. - * Combined with {@link HttpAgentCardLoader#normalize} pinning AgentCard - * URLs to the validated agent URI, this closes the SSRF bypass that would - * otherwise exist in the default {@code JdkA2AHttpClient} (which uses - * {@code Redirect.NORMAL}). + * SSRF-safe HTTP clients used by the A2A JSON-RPC transport, one per + * distinct response cap, reference-counted by the {@link Client}s using + * them. + * + *

    Each is backed by a {@link java.net.http.HttpClient} built with + * {@code followRedirects(NEVER)}, the configured connect timeout, and a + * streaming response-size cap, so the transport can neither follow HTTP + * redirects to internal addresses nor be made to buffer an oversized + * agent response. Combined with {@link HttpAgentCardLoader#normalize} + * pinning AgentCard URLs to the validated agent URI, this closes the + * SSRF bypass that would otherwise exist in the default + * {@code JdkA2AHttpClient} (which uses {@code Redirect.NORMAL}). + * + *

    A JDK HTTP client owns a selector thread and an executor, and + * neither {@code JdkA2AHttpClient} nor {@code Client.close()} releases + * it — so the factory tracks ownership itself. Clients built for the + * same cap share one (the overwhelmingly common case: every agent on + * the default), and the last one evicted takes it down. The reference + * count, not a fixed pool bound, is what keeps a caller who varies the + * cap per call from accumulating HTTP clients. */ - private final org.a2aproject.sdk.client.http.A2AHttpClient safeHttpClient; + private final Map pool = new HashMap<>(); + private final Map capByClient = new IdentityHashMap<>(); + private final Object poolLock = new Object(); DefaultClientFactory(AdcpHttpClient adcpHttpClient) { - this.safeHttpClient = new JdkA2AHttpClient( - adcpHttpClient.newHttpClientBuilder().build()); + this.adcpHttpClient = Objects.requireNonNull(adcpHttpClient, "adcpHttpClient"); } @Override - public Client create(AgentCard agentCard) throws A2AClientException { + public Client create(AgentCard agentCard, long maxResponseBytes) + throws A2AClientException { ClientConfig config = ClientConfig.builder() .setStreaming(true) .setUseClientPreference(true) .build(); - return Client.builder(agentCard) - .clientConfig(config) - .withTransport(JSONRPCTransport.class, - new JSONRPCTransportConfigBuilder().httpClient(safeHttpClient)) - .build(); + java.net.http.HttpClient httpClient = acquire(maxResponseBytes); + boolean handedOff = false; + try { + Client client = Client.builder(agentCard) + .clientConfig(config) + .withTransport(JSONRPCTransport.class, + new JSONRPCTransportConfigBuilder() + .httpClient(new JdkA2AHttpClient(httpClient))) + .build(); + synchronized (poolLock) { + capByClient.put(client, maxResponseBytes); + } + handedOff = true; + return client; + } finally { + if (!handedOff) { + release(maxResponseBytes); + } + } + } + + @Override + public void release(Client client) { + Long cap; + synchronized (poolLock) { + cap = capByClient.remove(client); + } + if (cap != null) { + release(cap.longValue()); + } + } + + @Override + public void close() { + List orphaned; + synchronized (poolLock) { + orphaned = pool.values().stream().map(PooledHttpClient::httpClient).toList(); + pool.clear(); + capByClient.clear(); + } + orphaned.forEach(DefaultClientFactory::shutdownQuietly); + } + + /** Visible for testing: how many pooled HTTP clients are still alive. */ + int pooledClientCount() { + synchronized (poolLock) { + return pool.size(); + } + } + + private java.net.http.HttpClient acquire(long maxResponseBytes) { + synchronized (poolLock) { + PooledHttpClient pooled = pool.get(maxResponseBytes); + if (pooled == null) { + pooled = new PooledHttpClient( + adcpHttpClient.newHttpClientBuilder(maxResponseBytes).build()); + pool.put(maxResponseBytes, pooled); + } + pooled.references++; + return pooled.httpClient(); + } + } + + private void release(long maxResponseBytes) { + java.net.http.HttpClient unused = null; + synchronized (poolLock) { + PooledHttpClient pooled = pool.get(maxResponseBytes); + if (pooled != null && --pooled.references <= 0) { + pool.remove(maxResponseBytes); + unused = pooled.httpClient(); + } + } + if (unused != null) { + shutdownQuietly(unused); + } + } + + /** + * {@code shutdownNow()} rather than {@code close()}: this can run under + * the connection cache lock during LRU eviction, and {@code close()} + * blocks until every in-flight exchange finishes. The owning A2A client + * has already been closed by then, which matches the existing eviction + * contract — an in-flight call on an evicted connection fails with an + * IOException and {@code ProtocolClient} retries. + */ + private static void shutdownQuietly(java.net.http.HttpClient httpClient) { + try { + httpClient.shutdownNow(); + } catch (RuntimeException e) { + log.debug("Error shutting down pooled A2A HTTP client: {}", + sanitizeLogText(e.getMessage())); + } + } + + private static final class PooledHttpClient { + private final java.net.http.HttpClient httpClient; + private int references; + + PooledHttpClient(java.net.http.HttpClient httpClient) { + this.httpClient = httpClient; + } + + java.net.http.HttpClient httpClient() { + return httpClient; + } } } @@ -345,14 +543,16 @@ private HttpAgentCardLoader(AdcpHttpClient adcpHttpClient, ObjectMapper objectMa } @Override - public AgentCard load(AgentConfig agent, Map headers) { + public AgentCard load(AgentConfig agent, Map headers, + long maxResponseBytes) { URI cardUri = buildAgentCardUri(agent.agentUri()); try { // Use getForAgentCard so that auth headers (e.g. Authorization) are forwarded // to private agent-card endpoints. Redirect-following is already blocked by // the underlying HttpClient (followRedirects=NEVER) and SSRF validation is // enforced inside getForAgentCard, so header leakage is not possible. - AdcpHttpResponse response = adcpHttpClient.getForAgentCard(cardUri, headers); + AdcpHttpResponse response = adcpHttpClient.getForAgentCard( + cardUri, headers, maxResponseBytes); if (response.statusCode() >= 200 && response.statusCode() < 300 && !response.truncated()) { AgentCard parsed = objectMapper.readValue(response.body(), AgentCard.class); return normalize(parsed, agent.agentUri()); diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/mcp/McpCaller.java b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/mcp/McpCaller.java index 1c3f98f..45c0e72 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/mcp/McpCaller.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/mcp/McpCaller.java @@ -5,6 +5,8 @@ import io.modelcontextprotocol.client.McpSyncClient; import io.modelcontextprotocol.spec.McpSchema; import org.adcontextprotocol.adcp.error.ProtocolError; +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; +import org.adcontextprotocol.adcp.http.AdcpHttpClient; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -16,14 +18,18 @@ * *

    Wraps the MCP SDK's {@code callTool()} API, extracting structured * content and deserializing to the target response type. + * + *

    The authoritative response size cap lives in the transport's HTTP client + * (see {@link org.adcontextprotocol.adcp.http.ResponseSizeLimitingHttpClient}), + * which aborts an oversized body while it streams. The content checks here are + * a backstop for payloads that arrive by some other route, and they translate a + * transport-level {@link ResponseTooLargeError} back into the typed error + * instead of burying it in a {@link ProtocolError}. */ public final class McpCaller { private static final Logger log = LoggerFactory.getLogger(McpCaller.class); - /** Maximum allowed TextContent length (10 MB, matching ObjectMapper limits). */ - private static final int MAX_CONTENT_LENGTH = 10 * 1024 * 1024; - private final ObjectMapper objectMapper; public McpCaller(ObjectMapper objectMapper) { @@ -41,6 +47,11 @@ public McpCaller(ObjectMapper objectMapper) { * @param client the connected MCP client * @param toolName the MCP tool name (e.g. "get_products") * @param args the merged arguments (including version envelope) + *

    Uses the SDK's default 4 MiB cap as its content backstop, not the cap + * configured on the calling client — the transport's own cap is the + * authoritative one. Prefer + * {@link #callTool(McpSyncClient, String, Map, Class, long)}. + * * @param responseType the expected response type * @param response type * @return the deserialized response @@ -48,21 +59,46 @@ public McpCaller(ObjectMapper objectMapper) { */ public T callTool(McpSyncClient client, String toolName, Map args, Class responseType) { + return callTool(client, toolName, args, responseType, + AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES); + } + + /** + * Calls an MCP tool and deserializes the response, rejecting content larger + * than {@code maxResponseBytes}. + * + * @param maxResponseBytes the effective response size cap for this call + * @throws ResponseTooLargeError if the agent's response exceeds the cap + */ + public T callTool(McpSyncClient client, String toolName, + Map args, Class responseType, + long maxResponseBytes) { try { McpSchema.CallToolRequest request = new McpSchema.CallToolRequest(toolName, args); McpSchema.CallToolResult result = client.callTool(request); - return extractResponse(result, responseType); + return extractResponse(result, responseType, maxResponseBytes); } catch (ProtocolError e) { throw e; } catch (Exception e) { + // The cap is enforced inside the JDK body subscriber, so the typed + // error arrives wrapped. Surface it rather than hiding it in a + // ProtocolError, which the retry path would treat as transient. + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(e); + if (tooLarge != null) { + throw tooLarge; + } + // The MCP SDK's exception messages carry server-supplied text; + // strip control characters before it reaches a log or an LLM. throw new ProtocolError("mcp", - "MCP callTool failed for " + toolName + ": " + e.getMessage(), e); + "MCP callTool failed for " + toolName + ": " + + sanitizeErrorText(e.getMessage()), e); } } @SuppressWarnings("unchecked") - private T extractResponse(McpSchema.CallToolResult result, Class responseType) { + private T extractResponse(McpSchema.CallToolResult result, Class responseType, + long maxResponseBytes) { // If the tool itself reported an error, surface it before trying // to deserialize the content as a success payload. if (Boolean.TRUE.equals(result.isError())) { @@ -101,11 +137,10 @@ private T extractResponse(McpSchema.CallToolResult result, Class response log.debug("Skipping TextContent with null text"); continue; } - if (text.length() > MAX_CONTENT_LENGTH) { - throw new ProtocolError("mcp", - "MCP response content exceeds size limit (" - + text.length() + " > " - + MAX_CONTENT_LENGTH + ")", null); + // UTF-8 is never fewer bytes than chars, so a char count over + // the cap is already over the cap — reject without measuring. + if (text.length() > maxResponseBytes) { + throw new ResponseTooLargeError(maxResponseBytes, text.length(), null); } try { return objectMapper.readValue(text, responseType); diff --git a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/mcp/McpConnectionManager.java b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/mcp/McpConnectionManager.java index 5e6e47b..a87abc9 100644 --- a/adcp/src/main/java/org/adcontextprotocol/adcp/transport/mcp/McpConnectionManager.java +++ b/adcp/src/main/java/org/adcontextprotocol/adcp/transport/mcp/McpConnectionManager.java @@ -10,6 +10,7 @@ import org.adcontextprotocol.adcp.auth.WwwAuthenticateParser; import org.adcontextprotocol.adcp.error.AuthenticationRequiredError; import org.adcontextprotocol.adcp.error.ProtocolError; +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; import org.adcontextprotocol.adcp.http.AdcpHttpClient; import org.adcontextprotocol.adcp.http.AdcpHttpResponse; import org.jspecify.annotations.Nullable; @@ -30,7 +31,10 @@ /** * Manages cached MCP client connections with LRU eviction. * - *

    Cache key: {@code agentUrl::tokenHash}. Max {@value #MAX_CACHE_SIZE} entries. + *

    Cache key: {@code agentUrl::tokenHash::maxResponseBytes}. Max + * {@value #MAX_CACHE_SIZE} entries. The response cap is part of the key because + * it is baked into the transport's {@link java.net.http.HttpClient}: a call that + * overrides the cap must not reuse a connection built with a different one. * Implements StreamableHTTP → SSE fallback per TS SDK behavior. * *

    Thread-safe: cache reads/writes use {@code cacheLock} (short-held, never @@ -104,10 +108,29 @@ public McpConnectionManager(Duration connectTimeout, Duration requestTimeout, */ public McpSyncClient getOrConnect(URI agentUri, Map headers, String tokenHash) { + return getOrConnect(agentUri, headers, tokenHash, null); + } + + /** + * Gets or creates a cached MCP client connection with an explicit response + * size cap. + * + * @param agentUri the agent's base URI + * @param headers auth + extra headers to inject into MCP requests + * @param tokenHash hash of the auth token (for cache keying) + * @param maxResponseBytes per-connection response cap, or {@code null} to + * use the client default + * ({@link AdcpHttpClient#maxCallResponseBytes()}) + * @return a connected {@link McpSyncClient} + */ + public McpSyncClient getOrConnect(URI agentUri, Map headers, + String tokenHash, + @Nullable Long maxResponseBytes) { if (closed) { throw new IllegalStateException("McpConnectionManager is closed"); } - String cacheKey = agentUri + "::" + tokenHash; + long cap = effectiveMaxResponseBytes(maxResponseBytes); + String cacheKey = cacheKey(agentUri, tokenHash, cap); // Fast path: check cache under short lock (no I/O) cacheLock.lock(); @@ -147,7 +170,7 @@ public McpSyncClient getOrConnect(URI agentUri, Map headers, } // Network I/O happens here — only blocks threads in the same stripe - McpSyncClient client = connectWithFallback(agentUri, headers, cacheKey); + McpSyncClient client = connectWithFallback(agentUri, headers, cacheKey, cap); cacheLock.lock(); try { @@ -167,22 +190,47 @@ public McpSyncClient getOrConnect(URI agentUri, Map headers, } /** - * Evicts a specific connection from the cache. + * Evicts every cached connection for the given agent URI and token hash, + * regardless of the response cap it was built with. */ public void evict(URI agentUri, String tokenHash) { - String cacheKey = agentUri + "::" + tokenHash; + String prefix = agentUri + "::" + tokenHash + "::"; cacheLock.lock(); try { - McpSyncClient evicted = cache.remove(cacheKey); - if (evicted != null) { - knownStreamableKeys.remove(cacheKey); - closeQuietly(evicted); + var it = cache.entrySet().iterator(); + while (it.hasNext()) { + var entry = it.next(); + if (entry.getKey().startsWith(prefix)) { + it.remove(); + knownStreamableKeys.remove(entry.getKey()); + closeQuietly(entry.getValue()); + } } } finally { cacheLock.unlock(); } } + /** The response cap applied when a call does not override it. */ + public long defaultMaxResponseBytes() { + return adcpHttpClient.maxCallResponseBytes(); + } + + private long effectiveMaxResponseBytes(@Nullable Long override) { + if (override == null) { + return adcpHttpClient.maxCallResponseBytes(); + } + if (override <= 0) { + throw new IllegalArgumentException( + "maxResponseBytes must be positive: " + override); + } + return override; + } + + private static String cacheKey(URI agentUri, String tokenHash, long maxResponseBytes) { + return agentUri + "::" + tokenHash + "::" + maxResponseBytes; + } + /** * Evicts all cached connections for the given agent URI, regardless of * token hash. Use this when auth credentials rotate so that stale @@ -241,17 +289,18 @@ private void evictOldest() { * legitimate 5xx errors and double-charged every cold connect. */ private McpSyncClient connectWithFallback(URI agentUri, Map headers, - String cacheKey) { + String cacheKey, long maxResponseBytes) { String url = agentUri.toString(); Map safe = sanitizeHeaders(headers); // Known-good StreamableHTTP endpoints skip the probe if (knownStreamableKeys.contains(cacheKey)) { try { - McpSyncClient client = buildAndInit(url, safe, true); + McpSyncClient client = buildAndInit(url, safe, true, maxResponseBytes); log.debug("Reconnected to {} via StreamableHTTP (cached)", agentUri); return client; } catch (Exception e) { + rethrowIfResponseTooLarge(e); if (isAuthError(e)) { throw probeAndBuildAuthError(agentUri, e); } @@ -266,7 +315,7 @@ private McpSyncClient connectWithFallback(URI agentUri, Map head boolean useStreamable = probeSupportsStreamableHttp(agentUri, safe); try { - McpSyncClient client = buildAndInit(url, safe, useStreamable); + McpSyncClient client = buildAndInit(url, safe, useStreamable, maxResponseBytes); if (useStreamable) { knownStreamableKeys.add(cacheKey); } @@ -274,6 +323,7 @@ private McpSyncClient connectWithFallback(URI agentUri, Map head useStreamable ? "StreamableHTTP" : "SSE"); return client; } catch (Exception e) { + rethrowIfResponseTooLarge(e); if (isAuthError(e)) { throw probeAndBuildAuthError(agentUri, e); } @@ -282,10 +332,11 @@ private McpSyncClient connectWithFallback(URI agentUri, Map head log.debug("StreamableHTTP init failed despite probe, trying SSE for {}", agentUri); try { - McpSyncClient client = buildAndInit(url, safe, false); + McpSyncClient client = buildAndInit(url, safe, false, maxResponseBytes); log.debug("Connected to {} via SSE (fallback)", agentUri); return client; } catch (Exception e2) { + rethrowIfResponseTooLarge(e2); if (isAuthError(e2)) { throw probeAndBuildAuthError(agentUri, e2); } @@ -328,29 +379,31 @@ private boolean probeSupportsStreamableHttp(URI agentUri, Map he return false; } catch (InterruptedException e) { Thread.currentThread().interrupt(); - log.debug("StreamableHTTP probe interrupted for {}: {}", agentUri, e.getMessage()); + log.debug("StreamableHTTP probe interrupted for {}: {}", + agentUri, sanitizeLogText(e.getMessage())); return true; } catch (IOException e) { - log.debug("StreamableHTTP probe failed for {}: {}", agentUri, e.getMessage()); + log.debug("StreamableHTTP probe failed for {}: {}", + agentUri, sanitizeLogText(e.getMessage())); return true; } } private McpSyncClient buildAndInit(String url, Map headers, - boolean useStreamable) { + boolean useStreamable, long maxResponseBytes) { var reqBuilder = java.net.http.HttpRequest.newBuilder().timeout(requestTimeout); McpClientTransport transport = useStreamable ? HttpClientStreamableHttpTransport.builder(url) .connectTimeout(connectTimeout) .requestBuilder(reqBuilder) - .clientBuilder(adcpHttpClient.newMcpClientBuilder()) + .clientBuilder(adcpHttpClient.newMcpClientBuilder(maxResponseBytes)) .httpRequestCustomizer((rb, method, uri, body, ctx) -> headers.forEach(rb::header)) .build() : HttpClientSseClientTransport.builder(url) .connectTimeout(connectTimeout) .requestBuilder(reqBuilder) - .clientBuilder(adcpHttpClient.newMcpClientBuilder()) + .clientBuilder(adcpHttpClient.newMcpClientBuilder(maxResponseBytes)) .httpRequestCustomizer((rb, method, uri, body, ctx) -> headers.forEach(rb::header)) .build(); @@ -386,6 +439,19 @@ private static boolean hasCrlf(String s) { return s.indexOf('\r') >= 0 || s.indexOf('\n') >= 0; } + /** + * Strips control characters and truncates text that originated with the + * peer, so an agent cannot inject CR/LF into a log line or into a + * {@link ProtocolError} message an adopter may feed to an LLM. + */ + static String sanitizeLogText(@Nullable String raw) { + if (raw == null || raw.isBlank()) { + return "(no detail)"; + } + String truncated = raw.length() > 256 ? raw.substring(0, 256) + "..." : raw; + return truncated.replaceAll("[\\p{Cc}]", ""); + } + // NOTE: MCP SDK 1.1.2 does not expose HTTP response headers on errors. // We work around this by sending a HEAD probe to the agent URI to // retrieve the WWW-Authenticate challenge. When the MCP SDK adds @@ -409,10 +475,10 @@ private AuthenticationRequiredError probeAndBuildAuthError(URI agentUri, Excepti } catch (InterruptedException probeEx) { Thread.currentThread().interrupt(); log.debug("Auth challenge probe interrupted for {}: {}", - agentUri, probeEx.getMessage()); + agentUri, sanitizeLogText(probeEx.getMessage())); } catch (IOException probeEx) { log.debug("Auth challenge probe failed for {}: {}", - agentUri, probeEx.getMessage()); + agentUri, sanitizeLogText(probeEx.getMessage())); } return new AuthenticationRequiredError(agentUri, challenge, null, cause); } @@ -425,8 +491,24 @@ private AuthenticationRequiredError probeAndBuildAuthError(URI agentUri, Excepti return WwwAuthenticateParser.parse(wwwAuth); } + /** + * An oversized response is a hard failure, not a transport hiccup: rethrow + * it before the SSE fallback or the auth-probe path can mask it as a + * connectivity problem and retry against the same hostile peer. + */ + private static void rethrowIfResponseTooLarge(Exception e) { + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(e); + if (tooLarge != null) { + throw tooLarge; + } + } + + /** Depth bound so a cyclic cause chain cannot spin this loop forever. */ + private static final int MAX_CAUSE_DEPTH = 32; + private boolean isAuthError(Exception e) { - for (Throwable t = e; t != null; t = t.getCause()) { + Throwable t = e; + for (int depth = 0; t != null && depth < MAX_CAUSE_DEPTH; t = t.getCause(), depth++) { String msg = t.getMessage(); if (msg == null) continue; // Check MCP SDK's error type first @@ -434,6 +516,7 @@ private boolean isAuthError(Exception e) { if (isAuthMessage(msg)) return true; } if (isAuthMessage(msg)) return true; + if (t.getCause() == t) break; } return false; } @@ -453,7 +536,7 @@ private void closeQuietly(McpSyncClient client) { client.close(); } } catch (Exception e) { - log.debug("Error closing MCP client: {}", e.getMessage()); + log.debug("Error closing MCP client: {}", sanitizeLogText(e.getMessage())); } } } diff --git a/adcp/src/test/java/org/adcontextprotocol/adcp/AdcpClientTest.java b/adcp/src/test/java/org/adcontextprotocol/adcp/AdcpClientTest.java index 2877cce..7b1e931 100644 --- a/adcp/src/test/java/org/adcontextprotocol/adcp/AdcpClientTest.java +++ b/adcp/src/test/java/org/adcontextprotocol/adcp/AdcpClientTest.java @@ -124,4 +124,26 @@ void builder_rejects_cross_major_version() { .adcpVersion(new AdcpVersion(2, null)) .build()); } + + @Test + void builder_rejects_a_non_positive_maxResponseBytes() { + assertThrows(ConfigurationError.class, + () -> AdcpClient.builder() + .agent(AgentConfig.mcp("test", AGENT_URI)) + .maxResponseBytes(0)); + assertThrows(ConfigurationError.class, + () -> AdcpClient.builder() + .agent(AgentConfig.mcp("test", AGENT_URI)) + .maxResponseBytes(-1)); + } + + @Test + void builder_accepts_a_raised_maxResponseBytes() { + try (AdcpClient client = AdcpClient.builder() + .agent(AgentConfig.mcp("test", AGENT_URI)) + .maxResponseBytes(32L * 1024 * 1024) + .build()) { + assertNotNull(client); + } + } } diff --git a/adcp/src/test/java/org/adcontextprotocol/adcp/AgentConfigTest.java b/adcp/src/test/java/org/adcontextprotocol/adcp/AgentConfigTest.java index 10c6259..357d497 100644 --- a/adcp/src/test/java/org/adcontextprotocol/adcp/AgentConfigTest.java +++ b/adcp/src/test/java/org/adcontextprotocol/adcp/AgentConfigTest.java @@ -222,4 +222,35 @@ void rejects_extraHeaders_with_crlf_in_value() { .extraHeaders(Map.of("X-Key", "bad\nvalue")) .build()); } + + @Test + void maxResponseBytes_defaults_to_null_meaning_client_default() { + AgentConfig config = AgentConfig.builder().id("a1").agentUri(AGENT_URI).build(); + assertNull(config.maxResponseBytes()); + } + + @Test + void maxResponseBytes_overrides_the_client_default() { + AgentConfig config = AgentConfig.builder() + .id("a1") + .agentUri(AGENT_URI) + .maxResponseBytes(8L * 1024 * 1024) + .build(); + assertEquals(8L * 1024 * 1024, config.maxResponseBytes()); + assertTrue(config.toString().contains("maxResponseBytes=8388608")); + } + + @Test + void rejects_a_non_positive_maxResponseBytes() { + // The cap cannot be disabled — agent responses are untrusted input. + assertThrows(org.adcontextprotocol.adcp.error.ConfigurationError.class, + () -> AgentConfig.builder().id("a1").agentUri(AGENT_URI) + .maxResponseBytes(0)); + assertThrows(org.adcontextprotocol.adcp.error.ConfigurationError.class, + () -> AgentConfig.builder().id("a1").agentUri(AGENT_URI) + .maxResponseBytes(-1)); + assertThrows(org.adcontextprotocol.adcp.error.ConfigurationError.class, + () -> new AgentConfig("a1", AGENT_URI, Protocol.MCP, null, null, null, null, + null, null, null, java.util.Map.of(), -5L)); + } } diff --git a/adcp/src/test/java/org/adcontextprotocol/adcp/http/AdcpHttpClientBodyCapTest.java b/adcp/src/test/java/org/adcontextprotocol/adcp/http/AdcpHttpClientBodyCapTest.java new file mode 100644 index 0000000..f001a28 --- /dev/null +++ b/adcp/src/test/java/org/adcontextprotocol/adcp/http/AdcpHttpClientBodyCapTest.java @@ -0,0 +1,301 @@ +package org.adcontextprotocol.adcp.http; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Body-cap behaviour of {@link AdcpHttpClient} against a live loopback server. + * + *

    Two caps, two behaviours: the probe paths ({@code send}/{@code get}/ + * {@code post}) stop reading at the small probe cap and flag the result, + * because they only ever look at the status line and headers. The agent-card + * path parses what it gets, so it is held to the strict call cap and throws + * instead of handing back a half-read card. + */ +class AdcpHttpClientBodyCapTest { + + private HttpServer server; + private ExecutorService serverExecutor; + private String base; + + @BeforeEach + void startServer() throws IOException { + server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + serverExecutor = Executors.newCachedThreadPool(); + server.setExecutor(serverExecutor); + + server.createContext("/chunked", exchange -> { + int bytes = intParam(exchange, "bytes"); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, 0); + writeFiller(exchange, bytes); + }); + + server.createContext("/declared", exchange -> { + int bytes = intParam(exchange, "bytes"); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, bytes); + writeFiller(exchange, bytes); + }); + + server.createContext("/.well-known/agent-card.json", exchange -> { + byte[] body = ("{\"name\":\"a\",\"description\":\"d\",\"version\":\"1\"," + + "\"url\":\"http://127.0.0.1/\"}").getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + + // Sends headers, then one byte per `delayMs`, forever. + server.createContext("/drip", exchange -> { + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, 0); + try (OutputStream os = exchange.getResponseBody()) { + while (true) { + os.write('x'); + os.flush(); + Thread.sleep(50); + } + } catch (IOException | InterruptedException expectedWhenClientGivesUp) { + // the client hit its read deadline + } + }); + + // Sends headers and then nothing at all. + server.createContext("/stall", exchange -> { + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, 0); + try { + Thread.sleep(30_000); + } catch (InterruptedException expected) { + Thread.currentThread().interrupt(); + } + try { + exchange.getResponseBody().close(); + } catch (IOException ignored) { + // client already gone + } + }); + + server.createContext("/echo-accept-encoding", exchange -> { + byte[] body = String.valueOf( + exchange.getRequestHeaders().getFirst("Accept-Encoding")) + .getBytes(StandardCharsets.UTF_8); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + + server.start(); + base = "http://127.0.0.1:" + server.getAddress().getPort(); + } + + @AfterEach + void stopServer() { + server.stop(0); + serverExecutor.shutdownNow(); + } + + private AdcpHttpClient client() { + return AdcpHttpClient.builder() + .ssrfPolicy(SsrfPolicy.permissive()) + .maxResponseBytes(4096) + .maxCallResponseBytes(16 * 1024) + .build(); + } + + @Test + void probe_response_over_the_probe_cap_is_flagged_not_thrown() throws Exception { + try (AdcpHttpClient client = client()) { + AdcpHttpResponse response = + client.get(URI.create(base + "/chunked?bytes=20000"), Map.of()); + + assertTrue(response.truncated(), + "probe callers must be able to see that the body is incomplete"); + assertEquals(4096, response.body().length); + } + } + + @Test + void agent_card_over_the_call_cap_throws() { + try (AdcpHttpClient client = client()) { + ResponseTooLargeError error = assertThrows(ResponseTooLargeError.class, + () -> client.getForAgentCard( + URI.create(base + "/chunked?bytes=" + (64 * 1024)), Map.of())); + + assertEquals(16 * 1024, error.limit()); + assertTrue(error.bytesRead() > 16 * 1024); + } + } + + @Test + void agent_card_with_an_over_cap_content_length_is_rejected_before_reading() { + try (AdcpHttpClient client = client()) { + ResponseTooLargeError error = assertThrows(ResponseTooLargeError.class, + () -> client.getForAgentCard( + URI.create(base + "/declared?bytes=" + (64 * 1024)), Map.of())); + + assertEquals(16 * 1024, error.limit()); + assertEquals(64 * 1024, error.bytesRead(), + "the declared length should be reported verbatim — nothing was read"); + } + } + + @Test + void agent_card_under_the_call_cap_is_returned_whole() throws Exception { + try (AdcpHttpClient client = client()) { + AdcpHttpResponse response = client.getForAgentCard( + URI.create(base + "/.well-known/agent-card.json"), Map.of()); + + assertEquals(200, response.statusCode()); + assertFalse(response.truncated()); + assertTrue(response.bodyAsString().contains("\"name\":\"a\"")); + } + } + + @Test + void agent_card_between_the_probe_cap_and_the_call_cap_still_succeeds() throws Exception { + // A 4 KiB probe cap would have silently truncated this card and sent the + // A2A loader off to a synthetic fallback; the 16 KiB call cap must not. + try (AdcpHttpClient client = client()) { + AdcpHttpResponse response = client.getForAgentCard( + URI.create(base + "/declared?bytes=10000"), Map.of()); + + assertFalse(response.truncated()); + assertEquals(10000, response.body().length); + } + } + + @Test + void outbound_requests_force_identity_encoding() throws Exception { + try (AdcpHttpClient client = client()) { + URI uri = URI.create(base + "/echo-accept-encoding"); + + assertEquals("identity", + client.get(uri, Map.of("Accept-Encoding", "gzip")).bodyAsString()); + assertEquals("identity", + client.getForAgentCard(uri, Map.of("Accept-Encoding", "gzip")) + .bodyAsString()); + } + } + + @Test + void a_dripping_peer_cannot_hold_the_calling_thread_past_the_read_timeout() { + // HttpRequest.timeout() only covers the status line and headers. Without + // a body deadline, a peer that sends headers and then drips one byte at + // a time pins this thread and its socket for as long as it takes to + // reach the cap — at 4 MiB, weeks. + try (AdcpHttpClient client = AdcpHttpClient.builder() + .ssrfPolicy(SsrfPolicy.permissive()) + .readTimeout(Duration.ofMillis(600)) + .build()) { + URI uri = URI.create(base + "/drip"); + long start = System.nanoTime(); + + assertThrows(IOException.class, () -> client.getForAgentCard(uri, Map.of())); + + assertTrue(Duration.ofNanos(System.nanoTime() - start).toSeconds() < 20, + "the read must give up at the deadline, not run to the 4 MiB cap"); + } + } + + @Test + void a_stalled_peer_cannot_hold_the_calling_thread_either() { + // Nothing arrives at all, so an in-loop deadline check never runs — the + // watchdog has to close the stream out from under the blocked read. + try (AdcpHttpClient client = AdcpHttpClient.builder() + .ssrfPolicy(SsrfPolicy.permissive()) + .readTimeout(Duration.ofMillis(600)) + .build()) { + URI uri = URI.create(base + "/stall"); + long start = System.nanoTime(); + + assertThrows(IOException.class, () -> client.getForAgentCard(uri, Map.of())); + + assertTrue(Duration.ofNanos(System.nanoTime() - start).toSeconds() < 20, + "a stalled body read must not block forever"); + } + } + + @Test + void agent_card_cap_can_be_overridden_per_fetch() throws Exception { + // An agent configured for a bigger cap must get it on discovery too, + // not just on the call — otherwise its connection dies on the card. + try (AdcpHttpClient client = client()) { + URI uri = URI.create(base + "/declared?bytes=" + (32 * 1024)); + + assertThrows(ResponseTooLargeError.class, + () -> client.getForAgentCard(uri, Map.of())); + + AdcpHttpResponse response = client.getForAgentCard(uri, Map.of(), 64 * 1024); + assertEquals(32 * 1024, response.body().length); + } + } + + @Test + void agent_card_fetch_rejects_a_non_positive_cap() { + try (AdcpHttpClient client = client()) { + assertThrows(IllegalArgumentException.class, + () -> client.getForAgentCard(URI.create(base + "/chunked?bytes=1"), + Map.of(), 0)); + } + } + + @Test + void probe_and_call_caps_have_independent_defaults() { + try (AdcpHttpClient client = AdcpHttpClient.builder().build()) { + assertEquals(AdcpHttpClient.DEFAULT_MAX_RESPONSE_BYTES, client.maxResponseBytes()); + assertEquals(AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES, + client.maxCallResponseBytes()); + assertEquals(4 * 1024, client.maxResponseBytes()); + assertEquals(4 * 1024 * 1024, client.maxCallResponseBytes()); + } + } + + private static void writeFiller(HttpExchange exchange, int bytes) throws IOException { + byte[] body = new byte[bytes]; + Arrays.fill(body, (byte) 'x'); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } catch (IOException expectedWhenCapAborts) { + // the client stopped reading + } + } + + private static int intParam(HttpExchange exchange, String name) { + String query = exchange.getRequestURI().getQuery(); + if (query == null) { + return 0; + } + for (String pair : query.split("&")) { + String[] kv = pair.split("=", 2); + if (kv.length == 2 && kv[0].equals(name)) { + return Integer.parseInt(kv[1]); + } + } + return 0; + } +} diff --git a/adcp/src/test/java/org/adcontextprotocol/adcp/http/ResponseSizeLimitingHttpClientTest.java b/adcp/src/test/java/org/adcontextprotocol/adcp/http/ResponseSizeLimitingHttpClientTest.java new file mode 100644 index 0000000..9db901a --- /dev/null +++ b/adcp/src/test/java/org/adcontextprotocol/adcp/http/ResponseSizeLimitingHttpClientTest.java @@ -0,0 +1,499 @@ +package org.adcontextprotocol.adcp.http; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import org.adcontextprotocol.adcp.error.ResponseTooLargeError; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests the streaming response-size cap that protects every MCP and A2A call + * from a memory-DoS by an oversized agent response. + * + *

    Each test drives a real loopback HTTP server through a real JDK + * {@link HttpClient}, because the whole point of the cap is that it fires + * while the body streams in — a check that only runs after the body is + * buffered would pass a unit test and still OOM the caller. + * + *

    Remove the cap and these fail: the oversized-body cases return a body + * instead of throwing. + */ +class ResponseSizeLimitingHttpClientTest { + + private static final int CAP = 64 * 1024; + + private HttpServer server; + private ExecutorService serverExecutor; + private HttpClient raw; + private String base; + + @BeforeEach + void startServer() throws IOException { + server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + serverExecutor = Executors.newCachedThreadPool(); + server.setExecutor(serverExecutor); + + // Chunked body of `bytes` bytes, no Content-Length. + server.createContext("/chunked", exchange -> { + long bytes = longParam(exchange, "bytes", 0); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, 0); + writeBytes(exchange, bytes); + }); + + // Body of `bytes` bytes with a declared Content-Length. + server.createContext("/declared", exchange -> { + long bytes = longParam(exchange, "bytes", 0); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, bytes); + writeBytes(exchange, bytes); + }); + + // `events` SSE events of `size` payload bytes each. + server.createContext("/sse", exchange -> { + long events = longParam(exchange, "events", 1); + long size = longParam(exchange, "size", 16); + exchange.getResponseHeaders().add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, 0); + String eol = switch (stringParam(exchange, "eol", "lf")) { + case "cr" -> "\r"; + case "crlf" -> "\r\n"; + default -> "\n"; + }; + byte[] filler = new byte[(int) size]; + Arrays.fill(filler, (byte) 'x'); + try (OutputStream os = exchange.getResponseBody()) { + for (long i = 0; i < events; i++) { + os.write("data: ".getBytes(StandardCharsets.UTF_8)); + os.write(filler); + os.write((eol + eol).getBytes(StandardCharsets.UTF_8)); + os.flush(); + } + } catch (IOException expectedWhenCapAborts) { + // the client cancelled mid-stream + } + }); + + // Echoes method, body, a custom header and Accept-Encoding. + server.createContext("/echo-request", exchange -> { + String requestBody = new String(exchange.getRequestBody().readAllBytes(), + StandardCharsets.UTF_8); + byte[] body = (exchange.getRequestMethod() + + "|" + requestBody + + "|" + exchange.getRequestHeaders().getFirst("X-Custom") + + "|" + exchange.getRequestHeaders().getFirst("Accept-Encoding")) + .getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "text/plain"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + + // Echoes the request's Accept-Encoding header back as the body. + server.createContext("/accept-encoding", exchange -> { + byte[] body = String.valueOf( + exchange.getRequestHeaders().getFirst("Accept-Encoding")) + .getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "text/plain"); + exchange.sendResponseHeaders(200, body.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(body); + } + }); + + server.start(); + base = "http://127.0.0.1:" + server.getAddress().getPort(); + raw = HttpClient.newHttpClient(); + } + + @AfterEach + void stopServer() { + raw.close(); + server.stop(0); + serverExecutor.shutdownNow(); + } + + // -- the cap itself -- + + @Test + void over_cap_body_without_content_length_throws() { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + IOException thrown = assertThrows(IOException.class, + () -> get(client, "/chunked?bytes=" + (CAP * 4L))); + + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(thrown); + assertNotNull(tooLarge, "expected a ResponseTooLargeError in the cause chain"); + assertEquals(CAP, tooLarge.limit()); + assertTrue(tooLarge.bytesRead() > CAP, + "bytesRead should report the count at the moment of the abort"); + } + + @Test + void over_cap_content_length_is_rejected_before_the_body_is_read() { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + long declared = CAP * 4L; + + IOException thrown = assertThrows(IOException.class, + () -> get(client, "/declared?bytes=" + declared)); + + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(thrown); + assertNotNull(tooLarge); + assertEquals(CAP, tooLarge.limit()); + // The declared length is reported verbatim — proof that nothing was read. + assertEquals(declared, tooLarge.bytesRead()); + } + + @Test + void under_cap_body_succeeds() throws Exception { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + HttpResponse response = get(client, "/chunked?bytes=" + (CAP / 2)); + + assertEquals(200, response.statusCode()); + assertEquals(CAP / 2, response.body().length()); + } + + @Test + void body_exactly_at_the_cap_succeeds() throws Exception { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + HttpResponse response = get(client, "/declared?bytes=" + CAP); + + assertEquals(CAP, response.body().length()); + } + + @Test + void per_call_override_raises_the_cap() throws Exception { + long body = CAP * 2L; + HttpClient capped = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + assertThrows(IOException.class, () -> get(capped, "/chunked?bytes=" + body)); + + HttpClient raised = ResponseSizeLimitingHttpClient.wrap(raw, body * 2); + assertEquals(body, get(raised, "/chunked?bytes=" + body).body().length()); + } + + @Test + void async_sends_are_capped_too() { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + CompletableFuture> future = client.sendAsync( + HttpRequest.newBuilder(URI.create(base + "/chunked?bytes=" + (CAP * 4L))).build(), + HttpResponse.BodyHandlers.ofString()); + + ExecutionException thrown = assertThrows(ExecutionException.class, + () -> future.get(30, TimeUnit.SECONDS)); + assertNotNull(ResponseTooLargeError.findIn(thrown)); + } + + // -- SSE -- + + @Test + void single_over_cap_sse_event_throws() { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + // One event, four times the cap. A GET SSE stream has no cumulative + // cap, so only the per-event cap can catch this. + IOException thrown = assertThrows(IOException.class, + () -> get(client, "/sse?events=1&size=" + (CAP * 4L))); + + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(thrown); + assertNotNull(tooLarge, "a single oversized SSE event must not escape the cap"); + assertEquals(CAP, tooLarge.limit()); + } + + @Test + void cumulative_sse_over_cap_throws_for_a_single_call() { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + // Eight events, each an eighth of the cap under it: no single event + // trips the per-event cap, but together they are over. + long size = CAP / 4; + IOException thrown = assertThrows(IOException.class, + () -> post(client, "/sse?events=8&size=" + size)); + + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(thrown); + assertNotNull(tooLarge, + "a call answered as a drip-feed of small SSE events must not escape the cap"); + assertEquals(CAP, tooLarge.limit()); + } + + @Test + void sse_content_type_does_not_bypass_the_cap_on_a_call() { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + // A peer answering a POST with text/event-stream instead of + // application/json gets no extra budget. + assertThrows(IOException.class, + () -> post(client, "/sse?events=1&size=" + (CAP * 4L))); + } + + @org.junit.jupiter.params.ParameterizedTest + @org.junit.jupiter.params.provider.ValueSource(strings = {"lf", "cr", "crlf"}) + void event_boundaries_reset_the_per_event_counter_for_every_sse_line_terminator(String eol) + throws Exception { + // SSE accepts CR, LF and CRLF as line terminators, and the JDK's + // fromLineSubscriber — which both transports sit behind — splits on all + // three. Missing one makes the per-event counter run on forever and + // kills a conforming peer's stream with a bogus ResponseTooLargeError. + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + HttpResponse response = + get(client, "/sse?events=8&size=" + (CAP / 4) + "&eol=" + eol); + + assertEquals(200, response.statusCode()); + assertTrue(response.body().length() > CAP, + "8 events of a quarter-cap each is twice the cap in total"); + } + + @org.junit.jupiter.params.ParameterizedTest + @org.junit.jupiter.params.provider.ValueSource(strings = {"lf", "cr", "crlf"}) + void a_single_over_cap_event_is_caught_for_every_line_terminator(String eol) { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + IOException thrown = assertThrows(IOException.class, + () -> get(client, "/sse?events=1&size=" + (CAP * 4L) + "&eol=" + eol)); + + assertNotNull(ResponseTooLargeError.findIn(thrown)); + } + + @Test + void long_lived_get_sse_stream_still_has_a_finite_cumulative_ceiling() { + // The listening-stream exemption is triggered by the peer's own + // Content-Type, so it must not mean "unbounded": a peer answering a GET + // with text/event-stream and a forever-stream of small events still + // gets torn down. + long tinyCap = 2048; + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, tinyCap); + long ceiling = tinyCap * ResponseSizeLimitingHttpClient.LISTENING_STREAM_MULTIPLIER; + long events = (ceiling / 512) + 16; + + IOException thrown = assertThrows(IOException.class, + () -> get(client, "/sse?events=" + events + "&size=500")); + + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(thrown); + assertNotNull(tooLarge, "a listening stream must still have a finite ceiling"); + assertEquals(ceiling, tooLarge.limit()); + } + + @Test + void long_lived_get_sse_stream_is_capped_per_event_not_cumulatively() throws Exception { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + // The MCP SSE session channel is a GET that stays open for the life of + // the connection and carries every server-to-client message. Capping + // its running total would kill healthy long sessions, so only the + // per-event cap applies: 8 events of a quarter-cap each (twice the cap + // in total) must stream through. + HttpResponse response = get(client, "/sse?events=8&size=" + (CAP / 4)); + + assertEquals(200, response.statusCode()); + assertTrue(response.body().length() > CAP, + "the listening stream should be allowed past the cumulative cap"); + assertEquals(8, response.body().split("\n\n", -1).length - 1); + } + + // -- encoding -- + + @Test + void outbound_requests_ask_for_identity_encoding() throws Exception { + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + // Counting compressed bytes would let a peer hide a huge payload behind + // a small gzip stream, so a caller-supplied Accept-Encoding is replaced. + HttpResponse response = client.send( + HttpRequest.newBuilder(URI.create(base + "/accept-encoding")) + .header("Accept-Encoding", "gzip, deflate") + .build(), + HttpResponse.BodyHandlers.ofString()); + + assertEquals("identity", response.body()); + } + + // -- plumbing -- + + @Test + void rewriting_accept_encoding_preserves_the_request_body_and_method() throws Exception { + // The MCP and A2A transports POST JSON-RPC payloads; the Accept-Encoding + // rewrite rebuilds the request, and dropping the body there would break + // every call while the cap tests still passed. + HttpClient client = ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + HttpResponse response = client.send( + HttpRequest.newBuilder(URI.create(base + "/echo-request")) + .header("Accept-Encoding", "gzip") + .header("X-Custom", "kept") + .POST(HttpRequest.BodyPublishers.ofString("{\"jsonrpc\":\"2.0\"}")) + .build(), + HttpResponse.BodyHandlers.ofString()); + + assertEquals("POST|{\"jsonrpc\":\"2.0\"}|kept|identity", response.body()); + } + + @Test + void wrap_rejects_a_non_positive_cap() { + assertThrows(IllegalArgumentException.class, + () -> ResponseSizeLimitingHttpClient.wrap(raw, 0)); + assertThrows(IllegalArgumentException.class, + () -> ResponseSizeLimitingHttpClient.wrap(raw, -1)); + } + + @Test + void wrapper_exposes_the_delegate_configuration() { + ResponseSizeLimitingHttpClient client = + ResponseSizeLimitingHttpClient.wrap(raw, CAP); + + assertEquals(CAP, client.maxResponseBytes()); + assertEquals(raw.followRedirects(), client.followRedirects()); + assertEquals(raw.version(), client.version()); + assertEquals(raw.connectTimeout(), client.connectTimeout()); + } + + @Test + void adcp_http_client_builds_capped_transport_clients() { + AdcpHttpClient adcp = AdcpHttpClient.builder() + .maxCallResponseBytes(123_456) + .build(); + try (adcp) { + HttpClient mcp = adcp.newMcpClientBuilder().build(); + HttpClient a2a = adcp.newHttpClientBuilder().build(); + HttpClient override = adcp.newHttpClientBuilder(999).build(); + try (mcp; a2a; override) { + assertEquals(123_456, + ((ResponseSizeLimitingHttpClient) mcp).maxResponseBytes()); + assertEquals(123_456, + ((ResponseSizeLimitingHttpClient) a2a).maxResponseBytes()); + assertEquals(999, + ((ResponseSizeLimitingHttpClient) override).maxResponseBytes()); + assertEquals(HttpClient.Redirect.NEVER, mcp.followRedirects()); + } + } + } + + @Test + void adcp_http_client_rejects_a_non_positive_transport_cap() { + try (AdcpHttpClient adcp = AdcpHttpClient.builder().build()) { + assertThrows(IllegalArgumentException.class, + () -> adcp.newHttpClientBuilder(0)); + assertThrows(IllegalArgumentException.class, + () -> AdcpHttpClient.builder().maxCallResponseBytes(-1)); + } + } + + @Test + void error_message_redacts_credentials_in_the_url() { + URI withSecrets = URI.create("https://svc:s3cr3t@agent.example.com/mcp?api_key=abc123"); + ResponseTooLargeError error = new ResponseTooLargeError(10, 20, withSecrets); + + assertFalse(error.getMessage().contains("s3cr3t"), + "userinfo must not reach a message that gets logged and shown to users"); + assertFalse(error.getMessage().contains("abc123"), + "query strings carry API keys"); + assertTrue(error.getMessage().contains("agent.example.com/mcp")); + assertEquals(withSecrets, error.url(), "url() still returns what was configured"); + } + + @Test + void find_in_walks_cause_and_suppressed_chains_without_looping() { + ResponseTooLargeError error = new ResponseTooLargeError(1, 2, null); + + assertNull(ResponseTooLargeError.findIn(null)); + assertEquals(error, ResponseTooLargeError.findIn(error)); + assertEquals(error, + ResponseTooLargeError.findIn(new IOException("wrapped", error))); + + IOException suppressing = new IOException("outer"); + suppressing.addSuppressed(error); + assertEquals(error, ResponseTooLargeError.findIn(suppressing)); + + // A cyclic cause chain must terminate rather than spin forever. + IOException first = new IOException("a"); + IOException second = new IOException("b"); + first.initCause(second); + second.initCause(first); + assertNull(ResponseTooLargeError.findIn(first)); + } + + // -- helpers -- + + private HttpResponse get(HttpClient client, String path) throws Exception { + return client.send(HttpRequest.newBuilder(URI.create(base + path)).build(), + HttpResponse.BodyHandlers.ofString()); + } + + private HttpResponse post(HttpClient client, String path) throws Exception { + return client.send( + HttpRequest.newBuilder(URI.create(base + path)) + .POST(HttpRequest.BodyPublishers.ofString("{}")) + .build(), + HttpResponse.BodyHandlers.ofString()); + } + + private static void writeBytes(HttpExchange exchange, long bytes) throws IOException { + byte[] chunk = new byte[8192]; + Arrays.fill(chunk, (byte) 'x'); + try (OutputStream os = exchange.getResponseBody()) { + long written = 0; + while (written < bytes) { + int n = (int) Math.min(chunk.length, bytes - written); + os.write(chunk, 0, n); + written += n; + } + } catch (IOException expectedWhenCapAborts) { + // the client cancelled mid-stream + } + } + + private static String stringParam(HttpExchange exchange, String name, String fallback) { + String query = exchange.getRequestURI().getQuery(); + if (query == null) { + return fallback; + } + for (String pair : query.split("&")) { + String[] kv = pair.split("=", 2); + if (kv.length == 2 && kv[0].equals(name)) { + return kv[1]; + } + } + return fallback; + } + + private static long longParam(HttpExchange exchange, String name, long fallback) { + String query = exchange.getRequestURI().getQuery(); + if (query == null) { + return fallback; + } + for (String pair : query.split("&")) { + List kv = Arrays.asList(pair.split("=", 2)); + if (kv.size() == 2 && kv.get(0).equals(name)) { + return Long.parseLong(kv.get(1)); + } + } + return fallback; + } +} diff --git a/adcp/src/test/java/org/adcontextprotocol/adcp/transport/a2a/A2aConnectionManagerTest.java b/adcp/src/test/java/org/adcontextprotocol/adcp/transport/a2a/A2aConnectionManagerTest.java index 76bb004..ffcfe25 100644 --- a/adcp/src/test/java/org/adcontextprotocol/adcp/transport/a2a/A2aConnectionManagerTest.java +++ b/adcp/src/test/java/org/adcontextprotocol/adcp/transport/a2a/A2aConnectionManagerTest.java @@ -37,11 +37,11 @@ void getOrConnect_reuses_cached_client_for_same_url() { AtomicInteger loaderCalls = new AtomicInteger(); AtomicInteger factoryCalls = new AtomicInteger(); manager = new A2aConnectionManager( - (agent, headers) -> { + (agent, headers, cap) -> { loaderCalls.incrementAndGet(); return testCard(agent.agentUri()); }, - agentCard -> { + (agentCard, maxResponseBytes) -> { factoryCalls.incrementAndGet(); return testClient(agentCard); }); @@ -65,11 +65,11 @@ void getOrConnect_different_auth_headers_get_separate_clients() { AtomicInteger factoryCalls = new AtomicInteger(); List> discoveryHeaders = new ArrayList<>(); manager = new A2aConnectionManager( - (agent, headers) -> { + (agent, headers, cap) -> { discoveryHeaders.add(Map.copyOf(headers)); return testCard(agent.agentUri()); }, - agentCard -> { + (agentCard, maxResponseBytes) -> { factoryCalls.incrementAndGet(); return testClient(agentCard); }); @@ -94,11 +94,11 @@ void getOrConnect_different_auth_headers_get_separate_clients() { void buildCacheKey_does_not_include_authorization_header() { List> discoveryHeaders = new ArrayList<>(); manager = new A2aConnectionManager( - (agent, headers) -> { + (agent, headers, cap) -> { discoveryHeaders.add(Map.copyOf(headers)); return testCard(agent.agentUri()); }, - A2aConnectionManagerTest::testClient); + (card, maxResponseBytes) -> testClient(card)); AgentConfig agent = AgentConfig.builder() .id("a2a-agent") @@ -125,8 +125,8 @@ void buildCacheKey_does_not_include_authorization_header() { void evict_exact_cache_hash_removes_only_matching_variant() { AtomicInteger factoryCalls = new AtomicInteger(); manager = new A2aConnectionManager( - (agent, headers) -> testCard(agent.agentUri()), - agentCard -> { + (agent, headers, cap) -> testCard(agent.agentUri()), + (agentCard, maxResponseBytes) -> { factoryCalls.incrementAndGet(); return testClient(agentCard); }); @@ -154,8 +154,8 @@ void evict_exact_cache_hash_removes_only_matching_variant() { void evict_forces_reconnect() { AtomicInteger factoryCalls = new AtomicInteger(); manager = new A2aConnectionManager( - (agent, headers) -> testCard(agent.agentUri()), - agentCard -> { + (agent, headers, cap) -> testCard(agent.agentUri()), + (agentCard, maxResponseBytes) -> { factoryCalls.incrementAndGet(); return testClient(agentCard); }); @@ -177,8 +177,8 @@ void evict_forces_reconnect() { @Test void getOrConnect_after_close_throws() { manager = new A2aConnectionManager( - (agent, headers) -> testCard(agent.agentUri()), - A2aConnectionManagerTest::testClient); + (agent, headers, cap) -> testCard(agent.agentUri()), + (card, maxResponseBytes) -> testClient(card)); manager.close(); AgentConfig agent = AgentConfig.builder() @@ -196,14 +196,14 @@ void httpAgentCardLoader_uses_agent_card_json_well_known_path() { List fetchedUris = new ArrayList<>(); // Use the package-private AgentCardLoader interface to capture the URI manager = new A2aConnectionManager( - (agent, headers) -> { + (agent, headers, cap) -> { // Reconstruct what HttpAgentCardLoader would build URI expected = URI.create(agent.agentUri().getScheme() + "://" + agent.agentUri().getAuthority() + "/.well-known/agent-card.json"); fetchedUris.add(expected); return testCard(agent.agentUri()); }, - A2aConnectionManagerTest::testClient); + (card, maxResponseBytes) -> testClient(card)); AgentConfig agent = AgentConfig.builder() .id("a2a-agent") @@ -230,11 +230,11 @@ void httpAgentCardLoader_buildAgentCardUri_uses_standard_path() { void authenticated_card_discovery_passes_auth_headers_to_loader() { List> capturedHeaders = new ArrayList<>(); manager = new A2aConnectionManager( - (agent, headers) -> { + (agent, headers, cap) -> { capturedHeaders.add(Map.copyOf(headers)); return testCard(agent.agentUri()); }, - A2aConnectionManagerTest::testClient); + (card, maxResponseBytes) -> testClient(card)); AgentConfig agent = AgentConfig.builder() .id("a2a-agent") diff --git a/adcp/src/test/java/org/adcontextprotocol/adcp/transport/a2a/ResponseSizeCapWiringTest.java b/adcp/src/test/java/org/adcontextprotocol/adcp/transport/a2a/ResponseSizeCapWiringTest.java new file mode 100644 index 0000000..7d48427 --- /dev/null +++ b/adcp/src/test/java/org/adcontextprotocol/adcp/transport/a2a/ResponseSizeCapWiringTest.java @@ -0,0 +1,225 @@ +package org.adcontextprotocol.adcp.transport.a2a; + +import org.a2aproject.sdk.client.Client; +import org.a2aproject.sdk.client.config.ClientConfig; +import org.a2aproject.sdk.client.transport.jsonrpc.JSONRPCTransport; +import org.a2aproject.sdk.client.transport.jsonrpc.JSONRPCTransportConfigBuilder; +import org.a2aproject.sdk.spec.A2AClientException; +import org.a2aproject.sdk.spec.AgentCapabilities; +import org.a2aproject.sdk.spec.AgentCard; +import org.a2aproject.sdk.spec.AgentInterface; +import org.adcontextprotocol.adcp.AgentConfig; +import org.adcontextprotocol.adcp.Protocol; +import org.adcontextprotocol.adcp.http.AdcpHttpClient; +import org.adcontextprotocol.adcp.http.SsrfPolicy; +import org.adcontextprotocol.adcp.schema.AdcpObjectMapperFactory; +import org.adcontextprotocol.adcp.transport.CallToolOptions; +import org.adcontextprotocol.adcp.transport.ProtocolClient; +import org.adcontextprotocol.adcp.transport.mcp.McpConnectionManager; +import org.junit.jupiter.api.Test; + +import java.net.URI; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Verifies that the response size cap reaches the transport that enforces it, + * and that the per-client → per-agent → per-call override order is honoured. + * + *

    The cap is baked into the transport's {@code HttpClient} at connect time, + * so "is the override wired" is answerable by recording the cap the client + * factory is handed. A regression here is silent — the call still works, it is + * just capped at the wrong number — which is exactly why it is asserted. + */ +class ResponseSizeCapWiringTest { + + /** An unroutable-but-fast-failing local port: the call fails, the cap is still recorded. */ + private static final URI AGENT_URI = URI.create("http://127.0.0.1:1/"); + + @Test + void client_default_cap_is_four_mebibytes() { + assertEquals(4L * 1024 * 1024, AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES); + try (AdcpHttpClient http = AdcpHttpClient.builder().build()) { + assertEquals(4L * 1024 * 1024, http.maxCallResponseBytes()); + } + } + + @Test + void default_cap_is_used_when_nothing_overrides_it() { + List caps = capsSeenBy(agent(null), CallToolOptions.DEFAULT); + + assertTrue(caps.size() >= 1, "the factory should have been asked for a client"); + assertEquals(AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES, caps.get(0)); + } + + @Test + void agent_config_overrides_the_client_default() { + List caps = capsSeenBy(agent(777_000L), CallToolOptions.DEFAULT); + + assertEquals(777_000L, caps.get(0)); + assertNotEquals(AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES, caps.get(0)); + } + + @Test + void call_options_override_the_agent_setting() { + List caps = capsSeenBy(agent(777_000L), + CallToolOptions.builder().maxResponseBytes(31_337).build()); + + assertEquals(31_337L, caps.get(0)); + } + + @Test + void connections_are_cached_per_cap() { + RecordingFactory factory = new RecordingFactory(); + try (A2aConnectionManager manager = new A2aConnectionManager( + (cfg, headers, cap) -> card(cfg.agentUri()), factory)) { + + AgentConfig agent = agent(null); + Client first = manager.getOrConnect(agent, Map.of(), "hash", 1_000L); + Client second = manager.getOrConnect(agent, Map.of(), "hash", 1_000L); + Client third = manager.getOrConnect(agent, Map.of(), "hash", 2_000L); + + assertEquals(first, second, "same cap should reuse the cached client"); + assertNotEquals(first, third, + "a different cap must not reuse a client built for the old cap"); + assertEquals(List.of(1_000L, 2_000L), factory.caps); + } + } + + @Test + void evict_removes_every_cap_variant() { + RecordingFactory factory = new RecordingFactory(); + try (A2aConnectionManager manager = new A2aConnectionManager( + (cfg, headers, cap) -> card(cfg.agentUri()), factory)) { + + AgentConfig agent = agent(null); + manager.getOrConnect(agent, Map.of(), "hash", 1_000L); + manager.getOrConnect(agent, Map.of(), "hash", 2_000L); + manager.evict(agent.agentUri(), "hash"); + + manager.getOrConnect(agent, Map.of(), "hash", 1_000L); + assertEquals(List.of(1_000L, 2_000L, 1_000L), factory.caps, + "both cap variants should have been evicted and rebuilt on demand"); + } + } + + @Test + void transport_http_clients_are_released_when_their_connection_is_evicted() { + // Neither JdkA2AHttpClient nor Client.close() releases the JDK HTTP + // client underneath, and each one owns a selector thread. A caller that + // varies maxResponseBytes per call would otherwise accumulate them + // without bound. + AdcpHttpClient http = AdcpHttpClient.builder().build(); + A2aConnectionManager.DefaultClientFactory factory = + new A2aConnectionManager.DefaultClientFactory(http); + try (http; A2aConnectionManager manager = new A2aConnectionManager( + (cfg, headers, cap) -> card(cfg.agentUri()), factory)) { + + AgentConfig agent = agent(null); + for (long cap = 1_000; cap <= 20_000; cap += 1_000) { + manager.getOrConnect(agent, Map.of(), "hash", cap); + } + assertTrue(factory.pooledClientCount() > 0); + + manager.evict(agent.agentUri(), "hash"); + assertEquals(0, factory.pooledClientCount(), + "every cap variant's HTTP client should be released on eviction"); + } + } + + @Test + void transport_http_clients_are_shared_while_the_cap_is_shared() { + AdcpHttpClient http = AdcpHttpClient.builder().build(); + A2aConnectionManager.DefaultClientFactory factory = + new A2aConnectionManager.DefaultClientFactory(http); + try (http; A2aConnectionManager manager = new A2aConnectionManager( + (cfg, headers, cap) -> card(cfg.agentUri()), factory)) { + + // Five agents, all on the default cap: one HTTP client between them. + for (int i = 0; i < 5; i++) { + manager.getOrConnect(agent(null), Map.of(), "hash-" + i, null); + } + assertEquals(1, factory.pooledClientCount()); + } + } + + @Test + void a_non_positive_override_is_rejected() { + RecordingFactory factory = new RecordingFactory(); + try (A2aConnectionManager manager = new A2aConnectionManager( + (cfg, headers, cap) -> card(cfg.agentUri()), factory)) { + AgentConfig agent = agent(null); + org.junit.jupiter.api.Assertions.assertThrows(IllegalArgumentException.class, + () -> manager.getOrConnect(agent, Map.of(), "hash", 0L)); + } + } + + // -- helpers -- + + /** + * Runs one {@code callTool} through a real {@link ProtocolClient} and + * returns the caps the client factory was handed. The call itself fails + * (nothing is listening) — the point is the cap that was resolved first. + */ + private static List capsSeenBy(AgentConfig agent, CallToolOptions options) { + RecordingFactory factory = new RecordingFactory(); + A2aConnectionManager a2a = new A2aConnectionManager( + (cfg, headers, cap) -> card(cfg.agentUri()), factory); + try (ProtocolClient client = new ProtocolClient( + AdcpObjectMapperFactory.create(), SsrfPolicy.permissive(), null, + new McpConnectionManager(Duration.ofSeconds(1)), a2a)) { + client.callTool(agent, "get_products", Map.of(), Object.class, options); + } catch (RuntimeException expected) { + // connection refused — we only care about the resolved cap + } + return factory.caps; + } + + private static AgentConfig agent(Long maxResponseBytes) { + AgentConfig.Builder builder = AgentConfig.builder() + .id("a2a-agent") + .agentUri(AGENT_URI) + .protocol(Protocol.A2A); + if (maxResponseBytes != null) { + builder.maxResponseBytes(maxResponseBytes); + } + return builder.build(); + } + + private static AgentCard card(URI uri) { + return AgentCard.builder() + .name("test") + .description("test agent") + .version("1.0") + .url(uri.toString()) + .preferredTransport("JSONRPC") + .capabilities(AgentCapabilities.builder() + .streaming(false).pushNotifications(false).build()) + .supportedInterfaces(List.of(new AgentInterface("JSONRPC", uri.toString()))) + .defaultInputModes(List.of("text")) + .defaultOutputModes(List.of("text")) + .skills(List.of()) + .build(); + } + + private static final class RecordingFactory implements A2aConnectionManager.ClientFactory { + private final List caps = new ArrayList<>(); + + @Override + public Client create(AgentCard agentCard, long maxResponseBytes) + throws A2AClientException { + caps.add(maxResponseBytes); + return Client.builder(agentCard) + .clientConfig(ClientConfig.builder() + .setStreaming(false).setUseClientPreference(true).build()) + .withTransport(JSONRPCTransport.class, new JSONRPCTransportConfigBuilder()) + .build(); + } + } +} diff --git a/adcp/src/test/java/org/adcontextprotocol/adcp/transport/mcp/McpConnectionManagerTest.java b/adcp/src/test/java/org/adcontextprotocol/adcp/transport/mcp/McpConnectionManagerTest.java index 0e57d30..06badba 100644 --- a/adcp/src/test/java/org/adcontextprotocol/adcp/transport/mcp/McpConnectionManagerTest.java +++ b/adcp/src/test/java/org/adcontextprotocol/adcp/transport/mcp/McpConnectionManagerTest.java @@ -49,6 +49,33 @@ void getOrConnect_after_close_throws() { () -> manager.getOrConnect(uri, java.util.Map.of(), "hash")); } + @Test + void default_response_cap_is_the_client_default() { + assertEquals( + org.adcontextprotocol.adcp.http.AdcpHttpClient.DEFAULT_MAX_CALL_RESPONSE_BYTES, + manager.defaultMaxResponseBytes()); + } + + @Test + void custom_http_client_cap_is_propagated() { + var http = org.adcontextprotocol.adcp.http.AdcpHttpClient.builder() + .maxCallResponseBytes(123_456) + .build(); + try (http; var mgr = new McpConnectionManager( + java.time.Duration.ofSeconds(1), java.time.Duration.ofSeconds(1), http)) { + assertEquals(123_456, mgr.defaultMaxResponseBytes()); + } + } + + @Test + void getOrConnect_rejects_a_non_positive_cap_override() { + var uri = java.net.URI.create("https://agent.example.com"); + assertThrows(IllegalArgumentException.class, + () -> manager.getOrConnect(uri, java.util.Map.of(), "hash", 0L)); + assertThrows(IllegalArgumentException.class, + () -> manager.getOrConnect(uri, java.util.Map.of(), "hash", -1L)); + } + @Test void double_close_is_safe() { manager.close(); diff --git a/docs/response-size-limits.md b/docs/response-size-limits.md new file mode 100644 index 0000000..ef11239 --- /dev/null +++ b/docs/response-size-limits.md @@ -0,0 +1,106 @@ +# Response size limits + +Every response an AdCP agent sends you is capped. The default is **4 MiB**, it is on without any +configuration, and it is enforced while the body streams in — an oversized response is aborted +mid-read rather than buffered and then rejected. + +Agent responses are untrusted input. Without the cap, one agent answering `get_products` with a +17.8 MB payload is enough to push a grader past 10 GB of memory and take down everything else +running in the same process. See [`specs/response-size-limits.md`](../specs/response-size-limits.md) +for the design, and [adcp#8190](https://github.com/adcontextprotocol/adcp/issues/8190) for the +cross-SDK contract. + +## What is capped + +MCP `callTool`, `tools/list` and discovery; A2A `message/send`; SSE streams on both; agent-card +fetches. Discovery probes — transport detection and `WWW-Authenticate` challenges, which read only +the status line and headers — keep a separate, smaller 4 KiB budget. + +## Raising the limit + +There is no opt-out. Every setter rejects a non-positive value: an agent that legitimately publishes +a very large catalogue gets a bigger number, not an uncapped read. + +```java +// Per client — applies to every agent this client talks to. +AdcpClient client = AdcpClient.builder() + .agent(agent) + .maxResponseBytes(32L * 1024 * 1024) + .build(); + +// Per agent — for the one seller with the enormous catalogue. +AgentConfig agent = AgentConfig.builder() + .id("big-catalogue-seller") + .agentUri(URI.create("https://agent.example.com")) + .maxResponseBytes(32L * 1024 * 1024) + .build(); + +// Per call — for the one call that needs it. +client.callTool("get_products", args, GetProductsResponse.class, + CallToolOptions.builder().maxResponseBytes(32L * 1024 * 1024).build()); +``` + +Resolution order is **call → agent → client**. The cap is baked into the transport connection, so +overriding it per call opens a separate pooled connection; prefer setting it on the `AgentConfig` +if an agent routinely needs more. + +## Handling the error + +An over-cap response throws `ResponseTooLargeError` (`code() == "RESPONSE_TOO_LARGE"`), carrying +`limit()`, `bytesRead()` and `url()`. Nothing is truncated — you never get back a half-read payload. + +The cap fires inside a JDK HTTP body subscriber, so by the time it reaches you the JDK or the +MCP/A2A transport may have wrapped it. Branch on it with `findIn`: + +```java +try { + return client.callTool("get_products", args, GetProductsResponse.class); +} catch (RuntimeException e) { + ResponseTooLargeError tooLarge = ResponseTooLargeError.findIn(e); + if (tooLarge != null) { + log.warn("{} answered with more than {} bytes — skipping", + agent.id(), tooLarge.limit()); + return null; + } + throw e; +} +``` + +`ResponseTooLargeError` is not retried by the SDK's evict-and-retry path: re-running the call would +just re-download the same oversized payload. + +## Upgrading + +This is a behaviour change. Before, the only limit was a 10 MB character check that ran *after* the +response had been fully buffered and parsed; now the effective limit is 4 MiB and it stops the read. +Two things to look at when you upgrade: + +- **Large catalogues.** If a call that used to return 6 MB now throws, raise the cap for that agent. + Start with the per-agent setter — it keeps the tighter default for everyone else. +- **Agent cards over 4 KiB.** Card fetches used to run on the 4 KiB probe budget, and a larger card + was silently truncated and replaced with a synthetic fallback. They now run on the 4 MiB call cap, + so a card with a realistic number of skills is read whole for the first time. If you were relying + on the fallback card's shape, you will now see the agent's real card. + +## SSE + +Streaming responses get no extra budget. Each SSE event is capped individually, and for anything +that answers a single call the running total across events is capped too — so a peer cannot evade +the limit by replying with `Content-Type: text/event-stream` or by drip-feeding small events. + +Note that on a capped call, progress and status events count against the same budget as the final +payload — a long-running task that emits many status updates can reach the limit without any single +message being large. + +The one exception is the long-lived MCP SSE session channel: a `GET` that stays open for the life of +the connection and carries every server-to-client message. Capping its running total at one call's +worth would kill a healthy long session after enough legitimate traffic, so it gets a cumulative +budget 256× larger, with the per-event cap unchanged. Per-event is the cap that bounds memory there +— the transport aggregates one event at a time and discards the bytes in between. + +## What the cap does not do + +It bounds what one peer can make you hold for one response. It is not a global memory budget: wire +bytes become a multiple of themselves once parsed into objects, and a batch job talking to many +agents multiplies that again. If you are fanning out — the shape that got OOM-killed in the first +place — bound your concurrency as well. diff --git a/specs/response-size-limits.md b/specs/response-size-limits.md new file mode 100644 index 0000000..094188d --- /dev/null +++ b/specs/response-size-limits.md @@ -0,0 +1,174 @@ +# Response size limits + +**Status:** Shipped — closes [#39](https://github.com/adcontextprotocol/adcp-sdk-java/issues/39) +**Tracks:** [`transport`](../ROADMAP.md#track-3--l0-transport-mcp--a2a) +**Decisions referenced:** D2 (JDK 21), D9 (MCP SDK pin), D10 (A2A SDK pin), D24 +**Cross-SDK tracking:** [adcontextprotocol/adcp#8190](https://github.com/adcontextprotocol/adcp/issues/8190) +**Adopter guide:** [`docs/response-size-limits.md`](../docs/response-size-limits.md) + +## Why this exists + +An AdCP agent's response is untrusted input. A hosted compliance grader was repeatedly OOM-killed +when a seller agent answered a single `get_products` call with a 17.8 MB payload: process memory +went past 10 GB, almost all of it outside the JS heap, and every agent graded in the same batch +starved. With a 4 MiB response cap the same run peaks at 1.6 GB. + +A buyer, crawler or grader built on an AdCP SDK should not be memory-DoS-able by one agent, and +that protection must not depend on every adopter remembering to opt in. So the cap is **on by +default**, in every SDK, at the same number. + +The Java-specific problem is that the SDK does not own the sockets on the call paths. MCP calls go +through `HttpClientStreamableHttpTransport` / `HttpClientSseClientTransport` and A2A calls through +`JdkA2AHttpClient` — all third-party, none of which exposes a body-size hook. What they *do* accept +is a JDK `java.net.http.HttpClient`. That is the choke point this design uses. + +## Defaults + +| Scope | Default | Where to change it | +|---|---|---| +| Agent responses (MCP `callTool`, A2A `message/send`, SSE, `tools/list`, agent card) | **4 MiB** (`4 * 1024 * 1024`) | `AdcpClient.Builder.maxResponseBytes(long)` | +| Per agent | inherits the client default | `AgentConfig.Builder.maxResponseBytes(long)` | +| Per call | inherits the agent default | `CallToolOptions.Builder.maxResponseBytes(long)` | +| Discovery probes (transport detection, `WWW-Authenticate` challenge) | 4 KiB | `AdcpHttpClient.Builder.maxResponseBytes(long)` | + +Resolution order for an agent response is **call → agent → client**. There is no opt-out: every +setter rejects a non-positive value. An adopter with a genuinely large catalogue raises the number; +nobody turns the cap off. + +```java +// Per client +AdcpClient client = AdcpClient.builder() + .agent(agent) + .maxResponseBytes(32L * 1024 * 1024) + .build(); + +// Per agent +AgentConfig agent = AgentConfig.builder() + .id("big-catalogue-seller") + .agentUri(URI.create("https://agent.example.com")) + .maxResponseBytes(32L * 1024 * 1024) + .build(); + +// Per call +client.callTool("get_products", args, GetProductsResponse.class, + CallToolOptions.builder().maxResponseBytes(32L * 1024 * 1024).build()); +``` + +## The error + +An over-cap response throws +[`ResponseTooLargeError`](../adcp/src/main/java/org/adcontextprotocol/adcp/error/ResponseTooLargeError.java) +(`code() == "RESPONSE_TOO_LARGE"`), carrying `limit()`, `bytesRead()` and `url()`. Nothing is ever +truncated on a path whose body is parsed. + +The cap fires inside a `java.net.http` body subscriber, so the JDK — and then the MCP/A2A +transports — may wrap it before it reaches the caller. Two things follow: + +- Callers who want to branch on it should use `ResponseTooLargeError.findIn(throwable)`, which walks + the cause and suppressed chains with a depth bound. +- The SDK's own transport code unwraps it (`McpCaller`, `A2aCaller`, `McpConnectionManager`, + `A2aConnectionManager`) so that `ProtocolClient`'s evict-and-retry path never re-downloads an + oversized payload as if it were a transient I/O blip. + +## How enforcement works + +[`ResponseSizeLimitingHttpClient`](../adcp/src/main/java/org/adcontextprotocol/adcp/http/ResponseSizeLimitingHttpClient.java) +decorates a JDK `HttpClient`. `AdcpHttpClient.newHttpClientBuilder(long)` returns a +`HttpClient.Builder` that builds one, so the MCP transports (`clientBuilder(...)`) and +`JdkA2AHttpClient` get a capped client without knowing it. Four properties: + +1. **Pre-reject on `Content-Length`.** A declared length over the cap fails before one body byte is + read. +2. **Count while reading.** Bytes are counted in `Flow.Subscriber.onNext`; the subscription is + cancelled the instant the cap is crossed. A check after the body is buffered or parsed would not + prevent the OOM this exists to stop. +3. **SSE per event and per call.** For `Content-Type: text/event-stream`, bytes are counted per + event (reset at each blank-line delimiter, accepting CR, LF and CRLF framing — all three are + legal SSE terminators and all three are what the JDK's `fromLineSubscriber` splits on) *and* + cumulatively for the whole response. Answering a POST with `text/event-stream` buys a peer no + extra budget, and neither does drip-feeding small events. +4. **Decompressed bytes.** Outbound requests are rewritten to `Accept-Encoding: identity`, + replacing any caller-supplied value, so a peer cannot hide a 1 GiB payload behind a small gzip + stream. (The JDK `HttpClient` never decompresses transparently, so wire bytes *are* the bytes the + caller holds.) +5. **A bounded read in time, not just in bytes.** `HttpRequest.timeout` covers only the status line + and headers — the JDK fires its timer on the response stage, not the body. So `getForAgentCard` + gives the body read its own deadline: a watchdog closes the stream at `readTimeout`, which + unblocks a stalled `read()`, and the read loop checks the same deadline so a drip-feed gets a + clear `HttpTimeoutException`. Without it, moving card fetches from the 4 KiB probe cap to the + 4 MiB call cap would have widened the slow-loris window by a factor of a thousand. + +### The one deliberate asymmetry: long-lived `GET` SSE streams + +The MCP SSE session channel is a `GET` that stays open for the life of the connection and carries +every server-to-client message. Holding its running total to a single call's cap would kill healthy +long sessions after enough legitimate traffic — so a `GET` answered with `text/event-stream` gets a +cumulative budget of `maxResponseBytes × 256` instead, with its per-event cap unchanged. + +Per-event is the cap that matters for memory here. OOM comes from how much is held at once, and the +transports aggregate one SSE event at a time, discarding bytes in between. Cumulative capping is +about bounding a *single call*, and a listening stream is not a call. Every path that *is* a call — +MCP StreamableHTTP POST, A2A `message/send`, A2A streaming POST — is capped cumulatively at +`maxResponseBytes` as well as per event. + +A response that declares a `Content-Length` is held to the per-call cap regardless of its +content type: a body of known, finite size is a payload, not a listening stream, so it gets no +share of the larger ceiling. + +The multiplier is finite on purpose. Which branch a response takes is decided partly by the +*peer's* `Content-Type`, so "no cumulative cap at all" would be a decision an attacker gets to make: +answer any `GET` with `text/event-stream` and a forever-stream of 1 KB events, and the stream never +ends. With the multiplier, the worst case is a long stream that eventually terminates. + +One consequence worth knowing when debugging: on a call that *is* capped cumulatively, progress and +status events count against the same budget as the final payload. A long-running `create_media_buy` +that emits many status updates can reach `RESPONSE_TOO_LARGE` without any single message being +large. + +### Probes keep their own, smaller cap + +`AdcpHttpClient.send` / `get` / `post` are discovery probes: a POST `ping` to detect StreamableHTTP +support, and HEAD/OPTIONS to read a `WWW-Authenticate` challenge. They use only the status line and +headers. Those keep the 4 KiB cap from `specs/ssrf-baseline.md` and stop reading rather than throw, +flagging the result via `AdcpHttpResponse.truncated()` — not silent, and no caller parses the body. + +`getForAgentCard` *is* parsed, so it moved to the strict 4 MiB call cap. Before this change it ran +on the 4 KiB probe cap and an agent card larger than 4 KiB — easily reached with a handful of skills +— was silently truncated and swapped for a synthetic fallback card. + +## What the cap does not bound + +The cap is **per in-flight response**. It bounds what one peer can make this process hold for one +call — it is not a global memory budget. Two things follow: + +- Wire bytes are not resident bytes. A 4 MiB body is buffered, copied out of the growing buffer, and + then parsed into an object graph; peak RSS for one response is a multiple of the cap. +- A caller fanning out across many agents multiplies that by its concurrency. The incident this cap + exists to prevent was a *batch* grader. Bound your concurrency too. + +## Connection caching + +The cap is baked into the transport's `HttpClient` at connect time, so it is part of the connection +cache key: `agentUrl::tokenHash::maxResponseBytes` (MCP) and `agentUri#cacheHash#maxResponseBytes` +(A2A). A call that overrides the cap gets a connection built for that cap rather than silently +reusing one built for another, and the override reaches agent-card discovery as well as the call +itself. Varying the cap per call therefore fragments the connection cache — prefer setting it once +on the `AgentConfig`. + +`A2aConnectionManager` keeps one JDK `HttpClient` per distinct cap, reference-counted by the A2A +clients using it. Neither `JdkA2AHttpClient` nor `Client.close()` releases that HTTP client, and +each one owns a selector thread, so the factory tracks ownership itself: clients on the same cap +share one (the common case — every agent on the default), and the last one evicted shuts it down. + +## Cross-SDK contract + +Behaviour that must stay identical across the TypeScript, Python, Go and Java SDKs +([adcp#8190](https://github.com/adcontextprotocol/adcp/issues/8190)): + +1. On by default for every agent response on the MCP and A2A call paths, including discovery and + tool listing. Default 4 MiB. Configurable per client, per agent and per call. +2. Oversized responses fail with a typed, documented error named `ResponseTooLarge…` per language + convention. Never a silent truncation. +3. Enforced while reading, with a `Content-Length` pre-reject. +4. SSE covered per event and cumulatively per call. +5. Decompressed bytes counted, or identity encoding forced. diff --git a/specs/ssrf-baseline.md b/specs/ssrf-baseline.md index de60ac9..f019409 100644 --- a/specs/ssrf-baseline.md +++ b/specs/ssrf-baseline.md @@ -19,7 +19,7 @@ The Java SDK has to match this baseline on every outbound discovery call — **n | Loopback service abuse | `http://localhost:9200/` Elasticsearch, `http://127.0.0.1:8500/v1/agent/services` Consul | | DNS rebinding | Domain that resolves to a public IP at TXT-record time and to `127.0.0.1` at connect time | | Redirect-to-internal | Public agent returns `Location: http://10.0.0.1/`; client follows | -| Slow-loris / huge-body | Public agent returns a 10 GB body to exhaust the prober's memory | +| Slow-loris / huge-body | Public agent returns a 10 GB body to exhaust the prober's memory (see [`response-size-limits.md`](response-size-limits.md)) | The first four classes target *what we connect to.* The fifth bypasses static URL validation by deferring the attack to redirect-follow time. The sixth weaponizes the prober's resource budget. @@ -28,7 +28,7 @@ The first four classes target *what we connect to.* The fifth bypasses static UR 1. **Resolve DNS once, validate the full address set.** For an outbound URL whose host is a name, resolve to the full address list, run every address through the address-guard, accept the call **only if all addresses pass**. Defense-in-depth against partial-block lists (host with one public + one private A record). 2. **Pin the connect to the first validated address.** Bypass `HttpClient`'s implicit re-resolution. Defense against DNS rebinding: even if the name's records change between resolution and connect, we connect to the IP we already validated. 3. **`redirect: 'manual'`.** No transparent redirect-follow. A 3xx from a probe is a result, not an action. Callers that need a follow re-run validation against the new URL. -4. **Body cap.** 4 KiB by default on probes (the agent-card / `tools/list` size class). Configurable per call site. +4. **Body cap.** 4 KiB by default on probes — transport detection and `WWW-Authenticate` challenges, where only the status line and headers are read. Configurable per call site. Paths whose body is *parsed* (agent cards, and every MCP/A2A call response) use the separate, strict 4 MiB cap from [`response-size-limits.md`](response-size-limits.md), which throws rather than truncating. ## Address guards (block list)