|
6 | 6 | import java.util.List; |
7 | 7 | import java.util.Map; |
8 | 8 | import java.util.ServiceLoader; |
| 9 | +import java.util.ServiceLoader.Provider; |
9 | 10 | import java.util.function.BiConsumer; |
10 | 11 | import java.util.function.Consumer; |
| 12 | +import java.util.stream.Collectors; |
11 | 13 |
|
12 | 14 | import io.a2a.client.config.ClientConfig; |
13 | 15 | import io.a2a.client.transport.spi.ClientTransport; |
14 | 16 | import io.a2a.client.transport.spi.ClientTransportConfig; |
15 | 17 | import io.a2a.client.transport.spi.ClientTransportConfigBuilder; |
16 | 18 | import io.a2a.client.transport.spi.ClientTransportProvider; |
| 19 | +import io.a2a.client.transport.spi.ClientTransportWrapper; |
17 | 20 | import io.a2a.spec.A2AClientException; |
18 | 21 | import io.a2a.spec.AgentCard; |
19 | 22 | import io.a2a.spec.AgentInterface; |
20 | 23 | import io.a2a.spec.TransportProtocol; |
21 | 24 | import org.jspecify.annotations.NonNull; |
22 | 25 | import org.jspecify.annotations.Nullable; |
23 | 26 |
|
| 27 | +import org.slf4j.Logger; |
| 28 | +import org.slf4j.LoggerFactory; |
| 29 | + |
24 | 30 | /** |
25 | 31 | * Builder for creating instances of {@link Client} to communicate with A2A agents. |
26 | 32 | * <p> |
@@ -96,6 +102,7 @@ public class ClientBuilder { |
96 | 102 |
|
97 | 103 | private static final Map<String, ClientTransportProvider<? extends ClientTransport, ? extends ClientTransportConfig<?>>> transportProviderRegistry = new HashMap<>(); |
98 | 104 | private static final Map<Class<? extends ClientTransport>, String> transportProtocolMapping = new HashMap<>(); |
| 105 | + private static final Logger LOGGER = LoggerFactory.getLogger(ClientBuilder.class); |
99 | 106 |
|
100 | 107 | static { |
101 | 108 | ServiceLoader<ClientTransportProvider> loader = ServiceLoader.load(ClientTransportProvider.class); |
@@ -318,7 +325,7 @@ private ClientTransport buildClientTransport() throws A2AClientException { |
318 | 325 | throw new A2AClientException("Missing required TransportConfig for " + agentInterface.protocolBinding()); |
319 | 326 | } |
320 | 327 |
|
321 | | - return clientTransportProvider.create(clientTransportConfig, agentCard, agentInterface); |
| 328 | + return wrap(clientTransportProvider.create(clientTransportConfig, agentCard, agentInterface), clientTransportConfig); |
322 | 329 | } |
323 | 330 |
|
324 | 331 | private Map<String, String> getServerPreferredTransports() throws A2AClientException { |
@@ -373,10 +380,55 @@ private AgentInterface findBestClientTransport() throws A2AClientException { |
373 | 380 | if (transportProtocol == null || transportUrl == null) { |
374 | 381 | throw new A2AClientException("No compatible transport found"); |
375 | 382 | } |
376 | | - if (! transportProviderRegistry.containsKey(transportProtocol)) { |
| 383 | + if (!transportProviderRegistry.containsKey(transportProtocol)) { |
377 | 384 | throw new A2AClientException("No client available for " + transportProtocol); |
378 | 385 | } |
379 | 386 |
|
380 | 387 | return new AgentInterface(transportProtocol, transportUrl); |
381 | 388 | } |
| 389 | + |
| 390 | + /** |
| 391 | + * Wraps the transport with all available transport wrappers discovered via ServiceLoader. |
| 392 | + * Wrappers are applied in reverse priority order (lowest priority first) to build a stack |
| 393 | + * where the highest priority wrapper is the outermost layer. |
| 394 | + * |
| 395 | + * @param transport the base transport to wrap |
| 396 | + * @param clientTransportConfig the transport configuration |
| 397 | + * @return the wrapped transport (or original if no wrappers are available/applicable) |
| 398 | + */ |
| 399 | + private ClientTransport wrap(ClientTransport transport, ClientTransportConfig<? extends ClientTransport> clientTransportConfig) { |
| 400 | + ServiceLoader<ClientTransportWrapper> wrapperLoader = ServiceLoader.load(ClientTransportWrapper.class); |
| 401 | + |
| 402 | + // Collect all wrappers, sort by priority, then reverse for stack application |
| 403 | + List<ClientTransportWrapper> wrappers = wrapperLoader.stream().map(Provider::get) |
| 404 | + .sorted() |
| 405 | + .collect(Collectors.toList()); |
| 406 | + |
| 407 | + if (wrappers.isEmpty()) { |
| 408 | + LOGGER.debug("No client transport wrappers found via ServiceLoader"); |
| 409 | + return transport; |
| 410 | + } |
| 411 | + LOGGER.debug(wrappers.size() + " client transport wrappers found via ServiceLoader"); |
| 412 | + |
| 413 | + // Reverse to apply lowest priority first (building stack with highest priority outermost) |
| 414 | + java.util.Collections.reverse(wrappers); |
| 415 | + |
| 416 | + // Apply wrappers to build stack |
| 417 | + ClientTransport wrapped = transport; |
| 418 | + for (ClientTransportWrapper wrapper : wrappers) { |
| 419 | + try { |
| 420 | + ClientTransport newWrapped = wrapper.wrap(wrapped, clientTransportConfig); |
| 421 | + if (newWrapped != wrapped) { |
| 422 | + LOGGER.debug("Applied transport wrapper: {} (priority: {})", |
| 423 | + wrapper.getClass().getName(), wrapper.priority()); |
| 424 | + } |
| 425 | + wrapped = newWrapped; |
| 426 | + } catch (Exception e) { |
| 427 | + LOGGER.warn("Failed to apply transport wrapper {}: {}", |
| 428 | + wrapper.getClass().getName(), e.getMessage(), e); |
| 429 | + } |
| 430 | + } |
| 431 | + |
| 432 | + return wrapped; |
| 433 | + } |
382 | 434 | } |
0 commit comments