Skip to content

Commit b6e6b59

Browse files
committed
Headers on the remote client transports: webSocketCustomizer and requestCustomizer
The JDK's HttpClient has no default headers, so neither WebSocketAcpClientTransport nor StreamableHttpAcpClientTransport could send an Authorization header, an API key or any other header an agent's endpoint requires; an application had to write its own transport. goose serve, for one, refuses every connection without X-Secret-Key. - WebSocketAcpClientTransport.webSocketCustomizer(Consumer<WebSocket.Builder>) runs on every connect attempt, after the connect timeout is set. A customizer that throws, or sets a header the JDK reserves for the handshake, fails that connect, which may be tried again. - StreamableHttpAcpClientTransport.requestCustomizer(Consumer<HttpRequest.Builder>) runs for every request the transport sends: the cleartext probe, initialize, each POST, every SSE stream opened or reopened, and the closing DELETE, so a token that expires is read again for each. It runs on a builder of its own and the result is copied with the protocol headers (Content-Type, Accept, Acp-Connection-Id, Acp-Session-Id) filtered out and the endpoint's URI restored, because a builder can replace a header but not remove one: the bootstrap initialize carries no connection id of the transport's own to replace a customizer's with. - Each HTTP request is now built inside its Mono, so a customizer that throws fails that Mono rather than escaping sendMessage, and a failed initialize may be sent again. initialize is built after the cleartext probe instead of being re-stamped with the settled version, so StreamableHttpRequests.pinned(HttpRequest) is gone. Verified against goose serve 1.52.0: without the header the WebSocket handshake is refused and initialize over HTTP answers 401; with X-Secret-Key from either customizer, initialize and session/new succeed.
1 parent 0d3e20c commit b6e6b59

6 files changed

Lines changed: 363 additions & 33 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
2828
Micronaut client configuration hard-coded the three as not advertised, so a Micronaut client
2929
could not advertise elicitation or boolean config options from its settings.
3030

31+
- **Headers on the remote client transports: `WebSocketAcpClientTransport.webSocketCustomizer(...)`
32+
and `StreamableHttpAcpClientTransport.requestCustomizer(...)`.** The JDK's `HttpClient` has no
33+
default headers, so neither transport could send an `Authorization` header, an API key or any
34+
other header an agent's endpoint requires; an application had to write its own transport.
35+
`goose serve`, for one, refuses every connection without its `X-Secret-Key`. The WebSocket
36+
customizer receives the `WebSocket.Builder` of each connect attempt; the HTTP one receives the
37+
`HttpRequest.Builder` of every request the transport sends (the cleartext probe, `initialize`,
38+
each POST, every SSE stream it opens or reopens, and the closing `DELETE`), so a token that
39+
expires is read again for each. The HTTP transport keeps the method, body, URI and its own
40+
headers (Content-Type, Accept, Acp-Connection-Id, Acp-Session-Id): a customizer's values for
41+
them are dropped. A customizer that throws, or sets a header the JDK restricts, fails that
42+
connect or that request's Mono rather than the caller, and a failed connect or `initialize` may
43+
be tried again. The HTTP transport now also builds each request inside its Mono, so a failure
44+
building one is reported there too.
45+
3146
- **Transport constructors without a JSON mapper:** `StreamableHttpAcpClientTransport(URI)`,
3247
`WebSocketAcpClientTransport(URI)`, `StreamableHttpAcpAgentTransport(int port, AcpAgentFactory)`
3348
and `StreamableHttpAcpServlet(AcpAgentFactory)` use `AcpJsonMapper.createDefault()`, as the stdio

‎acp-core/src/main/java/com/agentclientprotocol/sdk/client/transport/StreamableHttpAcpClientTransport.java‎

Lines changed: 42 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,16 @@
6161
* the server does not. The default client keeps cookies in a cookie manager of its own and
6262
* runs on a bounded pool of daemon threads; {@link StreamableHttpAcpClientTransportOptions}
6363
* sets its sizes and the number of SSE streams. Pass an {@link HttpClient} of your own for TLS,
64-
* proxy or authentication settings.
64+
* proxy or cookie settings.
65+
*
66+
* <p>An endpoint that requires authentication, such as one that expects an API key or a
67+
* bearer token in a header, gets it through {@link #requestCustomizer}, which every request
68+
* the transport sends passes through:
69+
*
70+
* <pre>{@code
71+
* var transport = new StreamableHttpAcpClientTransport(URI.create("https://agents.example.com/acp"))
72+
* .requestCustomizer(builder -> builder.header("Authorization", "Bearer " + tokens.current()));
73+
* }</pre>
6574
*
6675
* <p>{@link #closeGracefully()} closes the streams and sends {@code DELETE} for the connection,
6776
* waiting at most five seconds for the answer; {@link #close()} does the same and blocks for up
@@ -201,6 +210,31 @@ private static HttpClientBundle createDefaultHttpClient(StreamableHttpAcpClientT
201210
return HttpClientBundle.createDefault(options);
202211
}
203212

213+
/**
214+
* Customizes every HTTP request the transport sends, typically to add the headers the
215+
* endpoint requires: an {@code Authorization} header, an API key, a tenant. The JDK's
216+
* {@link HttpClient} has no default headers, so this is the only way to send one. It is
217+
* applied to each request as it is built, the cleartext probe, {@code initialize}, every
218+
* POST, every SSE stream (re)opened and the closing {@code DELETE}, so a header whose
219+
* value changes, such as a token that expires, is read again for each one. It runs on the
220+
* thread that sends the request; keep it quick, and refresh a token elsewhere.
221+
*
222+
* <p>The transport owns the method, the body, the URI and its own headers (Content-Type,
223+
* Accept, Acp-Connection-Id, Acp-Session-Id): values the customizer sets for any of them
224+
* are dropped or replaced, so a customizer cannot break the protocol by accident. Other
225+
* settings, such as a per-request timeout, are kept. The JDK refuses restricted headers
226+
* such as {@code Host} or {@code Connection}; setting one, or any exception the customizer
227+
* throws, fails that request's Mono. Call it before the first message is sent.
228+
* @param customizer applied to the builder of each request
229+
* @return this transport
230+
* @throws IllegalArgumentException if {@code customizer} is null
231+
*/
232+
public StreamableHttpAcpClientTransport requestCustomizer(Consumer<HttpRequest.Builder> customizer) {
233+
Assert.notNull(customizer, "The request customizer can not be null");
234+
this.requests.requestCustomizer(customizer);
235+
return this;
236+
}
237+
204238
/**
205239
* {@inheritDoc}
206240
* <p>It contacts nothing: it registers the handler and completes at once. The connection
@@ -249,17 +283,21 @@ private Mono<Void> initialize(AcpSchema.JSONRPCRequest request) {
249283
return Mono.error(new IllegalStateException("Transport is already initialized"));
250284
}
251285

252-
HttpRequest httpRequest;
286+
String json;
253287
try {
254-
httpRequest = requests.jsonPost(RouteScope.bootstrap(), jsonMapper.writeValueAsString(request));
288+
json = jsonMapper.writeValueAsString(request);
255289
}
256290
catch (IOException e) {
257291
initialized.set(false);
258292
return Mono.error(new AcpConnectionException("Failed to serialize initialize request", e));
259293
}
260294

295+
// Built after the probe, so it carries the HTTP version the probe settled on; built
296+
// inside defer, so a request customizer that throws fails this Mono rather than the
297+
// caller.
261298
return requests.upgradeCleartextToHttp2()
262-
.then(Mono.defer(() -> requests.sendAsync(requests.pinned(httpRequest), HttpResponse.BodyHandlers.ofString())))
299+
.then(Mono.defer(() -> requests.sendAsync(requests.jsonPost(RouteScope.bootstrap(), json),
300+
HttpResponse.BodyHandlers.ofString())))
263301
.flatMap(response -> readInitializeResponse(request, response))
264302
.flatMap(responseMessage -> streams.openConnectionStream().then(inbound.emit(responseMessage)))
265303
.doOnError(error -> {

‎acp-core/src/main/java/com/agentclientprotocol/sdk/client/transport/StreamableHttpRequests.java‎

Lines changed: 58 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,15 @@
1313
import java.nio.charset.StandardCharsets;
1414
import java.time.Duration;
1515
import java.util.Locale;
16+
import java.util.Set;
1617
import java.util.concurrent.ArrayBlockingQueue;
1718
import java.util.concurrent.CompletableFuture;
1819
import java.util.concurrent.ExecutorService;
1920
import java.util.concurrent.RejectedExecutionException;
2021
import java.util.concurrent.ThreadFactory;
2122
import java.util.concurrent.ThreadPoolExecutor;
2223
import java.util.concurrent.TimeUnit;
24+
import java.util.function.Consumer;
2325

2426
import com.agentclientprotocol.sdk.error.AcpConnectionException;
2527
import com.agentclientprotocol.sdk.util.AcpSchedulers;
@@ -30,8 +32,8 @@
3032

3133
/**
3234
* The HTTP exchanges of a Streamable HTTP client connection: the scope headers, the HTTP
33-
* version settled by the cleartext probe, and the POST, GET and DELETE requests with the
34-
* status and content type each must answer with. Completion signals are delivered on a
35+
* version settled by the cleartext probe, the application's request customizer, and the
36+
* POST, GET and DELETE requests with the status and content type each must answer with. Completion signals are delivered on a
3537
* bounded executor of their own, never on the HTTP client's.
3638
*/
3739
final class StreamableHttpRequests {
@@ -48,6 +50,14 @@ final class StreamableHttpRequests {
4850

4951
static final Duration PROBE_TIMEOUT = Duration.ofSeconds(5);
5052

53+
/**
54+
* The headers this transport owns, lower-cased. A request customizer's values for them are
55+
* dropped, including on the bootstrap {@code initialize}, which carries no connection id
56+
* of the transport's own for one set by the customizer to be replaced by.
57+
*/
58+
private static final Set<String> PROTOCOL_HEADERS = Set.of("content-type", "accept",
59+
HEADER_CONNECTION_ID.toLowerCase(Locale.ROOT), HEADER_SESSION_ID.toLowerCase(Locale.ROOT));
60+
5161
/** An HTTP client and the executor this transport created for it, if any. */
5262
record HttpClientBundle(HttpClient httpClient, @Nullable ExecutorService ownedExecutor) {
5363

@@ -75,6 +85,9 @@ static HttpClientBundle createDefault(StreamableHttpAcpClientTransportOptions op
7585

7686
private volatile @Nullable String connectionId;
7787

88+
private volatile Consumer<HttpRequest.Builder> requestCustomizer = builder -> {
89+
};
90+
7891
/**
7992
* HTTP version pinned for every request after the cleartext probe, or {@code null} to
8093
* use the client's own setting. Set to HTTP/1.1 when an {@code http://} server does not
@@ -106,6 +119,10 @@ static ThreadFactory daemonThreadFactory(String threadName) {
106119
};
107120
}
108121

122+
void requestCustomizer(Consumer<HttpRequest.Builder> requestCustomizer) {
123+
this.requestCustomizer = requestCustomizer;
124+
}
125+
109126
@Nullable String connectionId() {
110127
return connectionId;
111128
}
@@ -129,9 +146,10 @@ Mono<Void> upgradeCleartextToHttp2() {
129146
if (!"http".equalsIgnoreCase(endpointUri.getScheme()) || httpClient.version() != HttpClient.Version.HTTP_2) {
130147
return Mono.empty();
131148
}
132-
HttpRequest probe = HttpRequest.newBuilder(endpointUri).GET().build();
133-
// Cancelling the Mono on timeout cancels the HTTP exchange (see sendAsync).
134-
return sendAsync(probe, HttpResponse.BodyHandlers.discarding())
149+
// Customized like every other request: a gateway in front of the agent may refuse the
150+
// probe without the application's credentials. Cancelling the Mono on timeout cancels
151+
// the HTTP exchange (see sendAsync).
152+
return Mono.defer(() -> sendAsync(customized().GET().build(), HttpResponse.BodyHandlers.discarding()))
135153
.timeout(PROBE_TIMEOUT, AcpSchedulers.timeouts())
136154
.doOnNext(response -> {
137155
logger.debug("Cleartext probe to {} negotiated {}", endpointUri, response.version());
@@ -148,52 +166,66 @@ Mono<Void> upgradeCleartextToHttp2() {
148166
});
149167
}
150168

151-
/** Re-stamps a request built before the probe with the version the probe settled on. */
152-
HttpRequest pinned(HttpRequest request) {
153-
HttpClient.Version version = this.pinnedVersion;
154-
if (version == null) {
155-
return request;
156-
}
157-
return HttpRequest.newBuilder(request, (name, value) -> true).version(version).build();
158-
}
159-
169+
/** A customized builder for the endpoint, with the version the probe settled on. */
160170
private HttpRequest.Builder newRequest() {
161-
HttpRequest.Builder requestBuilder = HttpRequest.newBuilder(endpointUri);
171+
HttpRequest.Builder requestBuilder = customized();
162172
HttpClient.Version version = this.pinnedVersion;
163173
if (version != null) {
164174
requestBuilder.version(version);
165175
}
166176
return requestBuilder;
167177
}
168178

169-
/** A JSON POST of {@code json} in {@code scope}, not yet sent. */
179+
/**
180+
* A builder for the endpoint carrying what the application's customizer set, less the
181+
* protocol headers and with the endpoint's URI whatever the customizer did. The
182+
* customizer runs on a builder of its own and the result is copied, because a builder
183+
* can replace a header but never remove one. The transport sets the method, the body and
184+
* its own headers afterwards.
185+
*/
186+
private HttpRequest.Builder customized() {
187+
HttpRequest.Builder scratch = HttpRequest.newBuilder(endpointUri);
188+
requestCustomizer.accept(scratch);
189+
return HttpRequest
190+
.newBuilder(scratch.build(), (name, value) -> !PROTOCOL_HEADERS.contains(name.toLowerCase(Locale.ROOT)))
191+
.uri(endpointUri);
192+
}
193+
194+
/**
195+
* A JSON POST of {@code json} in {@code scope}, not yet sent. Builds the request, so it
196+
* throws whatever the request customizer throws; the Monos below build inside
197+
* {@code defer} and fail with it instead.
198+
*/
170199
HttpRequest jsonPost(RouteScope scope, String json) {
171-
HttpRequest.Builder builder = newRequest().header("Content-Type", CONTENT_TYPE_JSON)
172-
.header("Accept", CONTENT_TYPE_JSON);
200+
HttpRequest.Builder builder = newRequest().setHeader("Content-Type", CONTENT_TYPE_JSON)
201+
.setHeader("Accept", CONTENT_TYPE_JSON);
173202
addScopeHeaders(builder, scope);
174203
return builder.POST(HttpRequest.BodyPublishers.ofString(json, StandardCharsets.UTF_8)).build();
175204
}
176205

177206
/** Posts a message that the server must accept with 202. */
178207
Mono<Void> postAccepted(RouteScope scope, String json) {
179-
return sendAsync(jsonPost(scope, json), HttpResponse.BodyHandlers.discarding())
208+
return Mono.defer(() -> sendAsync(jsonPost(scope, json), HttpResponse.BodyHandlers.discarding()))
180209
.flatMap(response -> expectStatus(response, 202, "for POST"));
181210
}
182211

183212
/** Opens the SSE stream of {@code scope}; emits its body or an error, never completes empty. */
184213
Mono<InputStream> openEventStream(RouteScope scope) {
185-
HttpRequest.Builder builder = newRequest().GET().header("Accept", CONTENT_TYPE_EVENT_STREAM);
186-
addScopeHeaders(builder, scope);
187-
return sendAsync(builder.build(), HttpResponse.BodyHandlers.ofInputStream())
214+
return Mono.defer(() -> {
215+
HttpRequest.Builder builder = newRequest().GET().setHeader("Accept", CONTENT_TYPE_EVENT_STREAM);
216+
addScopeHeaders(builder, scope);
217+
return sendAsync(builder.build(), HttpResponse.BodyHandlers.ofInputStream());
218+
})
188219
.flatMap(response -> expectStatus(response, 200, "when opening SSE stream")
189220
.then(expectContentType(response, CONTENT_TYPE_EVENT_STREAM, "response"))
190221
.then(Mono.fromSupplier(response::body)));
191222
}
192223

193224
/** Deletes the connection; the server must accept with 202. */
194225
Mono<Void> deleteConnection(String id) {
195-
HttpRequest request = newRequest().DELETE().header(HEADER_CONNECTION_ID, id).build();
196-
return sendAsync(request, HttpResponse.BodyHandlers.discarding())
226+
return Mono
227+
.defer(() -> sendAsync(newRequest().DELETE().setHeader(HEADER_CONNECTION_ID, id).build(),
228+
HttpResponse.BodyHandlers.discarding()))
197229
.flatMap(response -> expectStatus(response, 202, "for DELETE"));
198230
}
199231

@@ -215,10 +247,10 @@ static Mono<Void> expectContentType(HttpResponse<?> response, String expected, S
215247

216248
private void addScopeHeaders(HttpRequest.Builder builder, RouteScope scope) {
217249
if (!scope.isBootstrap()) {
218-
builder.header(HEADER_CONNECTION_ID, requireConnectionId());
250+
builder.setHeader(HEADER_CONNECTION_ID, requireConnectionId());
219251
}
220252
if (scope.isSession()) {
221-
builder.header(HEADER_SESSION_ID, scope.boundSessionId());
253+
builder.setHeader(HEADER_SESSION_ID, scope.boundSessionId());
222254
}
223255
}
224256

‎acp-core/src/main/java/com/agentclientprotocol/sdk/client/transport/WebSocketAcpClientTransport.java‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,14 @@
5656
* skipped. When the agent closes the connection, or it fails, {@link #awaitTermination()} ends
5757
* and the client's pending requests fail.
5858
*
59+
* <p>An endpoint that requires authentication, such as one that expects an API key or a
60+
* bearer token in a header, gets it through {@link #webSocketCustomizer}:
61+
*
62+
* <pre>{@code
63+
* var transport = new WebSocketAcpClientTransport(URI.create("wss://agents.example.com/acp"))
64+
* .webSocketCustomizer(builder -> builder.header("Authorization", "Bearer " + token));
65+
* }</pre>
66+
*
5967
* <p>The transport is thread-safe: messages may be sent from any thread, and one daemon thread
6068
* of its own ({@code acp-ws-client-outbound}) sends them one frame at a time. The default
6169
* HTTP client runs on a pool of daemon threads named {@code acp-ws-client}.
@@ -108,6 +116,9 @@ public class WebSocketAcpClientTransport implements AcpClientTransport {
108116

109117
private Duration connectTimeout = Duration.ofSeconds(30);
110118

119+
private Consumer<WebSocket.Builder> webSocketCustomizer = builder -> {
120+
};
121+
111122
/**
112123
* Creates a transport for the WebSocket endpoint at {@code serverUri}, with an HTTP client
113124
* of its own and the default JSON mapper ({@link AcpJsonMapper#createDefault()}).
@@ -186,6 +197,26 @@ public WebSocketAcpClientTransport connectTimeout(Duration timeout) {
186197
return this;
187198
}
188199

200+
/**
201+
* Customizes the WebSocket handshake before it is sent, typically to add the headers the
202+
* endpoint requires: an {@code Authorization} header, an API key, a tenant. The JDK's
203+
* {@link java.net.http.HttpClient} has no default headers, so this is the only way to send
204+
* one. Runs on every {@link #connect} attempt, after the transport has set its connect
205+
* timeout, so the customizer may replace that too; a header whose value changes, such as
206+
* a token that expires, is read again by each attempt. The JDK refuses headers that belong
207+
* to the handshake itself ({@code Connection}, {@code Upgrade}, {@code Host},
208+
* {@code Sec-WebSocket-*}): setting one, or any exception the customizer throws, fails
209+
* that connect, which may then be tried again. Call it before connecting.
210+
* @param customizer applied to the builder of each handshake
211+
* @return this transport
212+
* @throws IllegalArgumentException if {@code customizer} is null
213+
*/
214+
public WebSocketAcpClientTransport webSocketCustomizer(Consumer<WebSocket.Builder> customizer) {
215+
Assert.notNull(customizer, "The WebSocket customizer can not be null");
216+
this.webSocketCustomizer = customizer;
217+
return this;
218+
}
219+
189220
/**
190221
* {@inheritDoc}
191222
* <p>Opens the WebSocket connection when the returned Mono is subscribed, and completes
@@ -204,9 +235,9 @@ public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
204235

205236
// Build WebSocket connection with listener; frames that arrive before the handshake
206237
// completes wait in the inbound sink.
207-
return httpClient.newWebSocketBuilder()
208-
.connectTimeout(connectTimeout)
209-
.buildAsync(serverUri, new AcpWebSocketListener());
238+
WebSocket.Builder builder = httpClient.newWebSocketBuilder().connectTimeout(connectTimeout);
239+
webSocketCustomizer.accept(builder);
240+
return builder.buildAsync(serverUri, new AcpWebSocketListener());
210241
}).doOnSuccess(ws -> {
211242
this.webSocket = ws;
212243
// Only an open connection takes the inbound sink's one subscriber, so a connect that

0 commit comments

Comments
 (0)