Skip to content

fix(transport): resolve SSE stream timeout in JdkHttpTransport #1322

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jujn wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(transport): resolve SSE stream timeout in JdkHttpTransport #1322

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 54 additions & 2 deletions agentscope-core/src/main/java/io/agentscope/core/model/transport/HttpTransportConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,21 @@ public class HttpTransportConfig {
/** Default connect timeout: 30 seconds. */
public static final Duration DEFAULT_CONNECT_TIMEOUT = Duration.ofSeconds(30);

/** Default read timeout: 5 minutes (for long-running model calls). */
/** Default response timeout (TTFT): 5 minutes (Time To First Token for streaming). */
public static final Duration DEFAULT_RESPONSE_TIMEOUT = Duration.ofMinutes(5);

/** Default stream idle timeout: 30 seconds (Maximum wait time between consecutive data chunks). */
public static final Duration DEFAULT_STREAM_IDLE_TIMEOUT = Duration.ofSeconds(30);

/** Default read timeout: 5 minutes (Overall timeout for non-streaming calls). */
public static final Duration DEFAULT_READ_TIMEOUT = Duration.ofMinutes(5);

/** Default write timeout: 30 seconds. */
public static final Duration DEFAULT_WRITE_TIMEOUT = Duration.ofSeconds(30);

private final Duration connectTimeout;
private final Duration responseTimeout;
private final Duration streamIdleTimeout;
private final Duration readTimeout;
private final Duration writeTimeout;
private final int maxIdleConnections;
Expand All @@ -45,6 +53,8 @@ public class HttpTransportConfig {

private HttpTransportConfig(Builder builder) {
this.connectTimeout = builder.connectTimeout;
this.responseTimeout = builder.responseTimeout;
this.streamIdleTimeout = builder.streamIdleTimeout;
this.readTimeout = builder.readTimeout;
this.writeTimeout = builder.writeTimeout;
this.maxIdleConnections = builder.maxIdleConnections;
Expand All @@ -64,7 +74,25 @@ public Duration getConnectTimeout() {
}

/**
* Get the read timeout.
* Get the response timeout (Time To First Token for streaming).
*
* @return the response timeout duration
*/
public Duration getResponseTimeout() {
return responseTimeout;
}

/**
* Get the stream idle timeout (maximum time between two consecutive data chunks).
*
* @return the stream idle timeout duration
*/
public Duration getStreamIdleTimeout() {
return streamIdleTimeout;
}

/**
* Get the read timeout(for non-streaming).
*
* @return the read timeout duration
*/
Expand Down Expand Up @@ -153,6 +181,8 @@ public static HttpTransportConfig defaults() {
*/
public static class Builder {
private Duration connectTimeout = DEFAULT_CONNECT_TIMEOUT;
private Duration responseTimeout = DEFAULT_RESPONSE_TIMEOUT;
private Duration streamIdleTimeout = DEFAULT_STREAM_IDLE_TIMEOUT;
private Duration readTimeout = DEFAULT_READ_TIMEOUT;
private Duration writeTimeout = DEFAULT_WRITE_TIMEOUT;
private int maxIdleConnections = 5;
Expand All @@ -172,6 +202,28 @@ public Builder connectTimeout(Duration connectTimeout) {
return this;
}

/**
* Set the response timeout (Time To First Byte).
*
* @param responseTimeout the response timeout duration
* @return this builder
*/
public Builder responseTimeout(Duration responseTimeout) {
this.responseTimeout = responseTimeout;
return this;
}

/**
* Set the stream idle timeout.
*
* @param streamIdleTimeout the stream idle timeout duration
* @return this builder
*/
public Builder streamIdleTimeout(Duration streamIdleTimeout) {
this.streamIdleTimeout = streamIdleTimeout;
return this;
}

/**
* Set the read timeout.
*
Expand Down
111 changes: 65 additions & 46 deletions agentscope-core/src/main/java/io/agentscope/core/model/transport/JdkHttpTransport.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
Expand Down Expand Up @@ -174,7 +174,7 @@ public HttpResponse execute(HttpRequest request) throws HttpTransportException {
throw new HttpTransportException("Transport has been closed");
}

var jdkRequest = buildJdkRequest(request);
var jdkRequest = buildJdkRequest(request, false);

try {
var response = client.send(jdkRequest, BodyHandlers.ofString());
Expand All @@ -193,50 +193,64 @@ public Flux<String> stream(HttpRequest request) {
return Flux.error(new HttpTransportException("Transport has been closed"));
}

var jdkRequest = buildJdkRequest(request);

// Check status code and read error body immediately when CompletableFuture completes
// to avoid stream being closed before we can read it
CompletableFuture<java.net.http.HttpResponse<InputStream>> future =
client.sendAsync(jdkRequest, BodyHandlers.ofInputStream())
.thenApply(
response -> {
int statusCode = response.statusCode();
if (statusCode < 200 || statusCode >= 300) {
// Read error body immediately while stream is still open
String errorBody = readInputStream(response.body());
log.warn(
"HTTP request failed. URL: {} | Status: {} | Error:"
+ " {}",
request.getUrl(),
var jdkRequest = buildJdkRequest(request, true);

// Use Mono.fromFuture() to ensure lazy execution and proper cancellation propagation.
// This prevents "ghost connections" from leaking if the downstream cancels or times out
// before headers arrive.
return Mono.fromFuture(() -> client.sendAsync(jdkRequest, BodyHandlers.ofInputStream()))
.flatMapMany(
response -> {
int statusCode = response.statusCode();
if (statusCode < 200 || statusCode >= 300) {
String errorBody = readInputStream(response.body());
log.warn(
"HTTP request failed. URL: {} | Status: {} | Error: {}",
request.getUrl(),
statusCode,
errorBody);
return Flux.error(
new HttpTransportException(
"HTTP request failed with status "
+ statusCode
+ " | "
+ errorBody,
statusCode,
errorBody);
throw new CompletionException(
new HttpTransportException(
"HTTP request failed with status "
+ statusCode
+ " | "
+ errorBody,
statusCode,
errorBody));
}
return response;
});

return Mono.fromCompletionStage(future)
.flatMapMany(response -> processStreamResponse(response, request))
.publishOn(Schedulers.boundedElastic())
errorBody));
}
return processStreamResponse(response, request);
})
.timeout(
// Timeout strategy 1: Time To First Token (TTFT).
// The maximum time to wait for the first piece of data.
Mono.delay(
config.getResponseTimeout() != null
? config.getResponseTimeout()
: HttpTransportConfig.DEFAULT_RESPONSE_TIMEOUT),

// Timeout strategy 2: Inter-token gap (Stream Idle Timeout).
// The maximum time to wait between receiving two consecutive data chunks.
data ->
Mono.delay(
config.getStreamIdleTimeout() != null
? config.getStreamIdleTimeout()
: HttpTransportConfig.DEFAULT_STREAM_IDLE_TIMEOUT))
.onErrorMap(
e -> !(e instanceof HttpTransportException),
e -> {
if (e instanceof TimeoutException) {
return new HttpTransportException(
"Stream timeout: " + e.getMessage(), e);
}
if (e instanceof HttpTransportException) {
return e;
}
Throwable cause = e instanceof CompletionException ? e.getCause() : e;
if (cause instanceof HttpTransportException) {
return (HttpTransportException) cause;
return cause;
}
return new HttpTransportException(
"SSE/NDJSON stream failed: " + e.getMessage(), e);
})
.subscribeOn(Schedulers.boundedElastic());
});
}

private Flux<String> processStreamResponse(
Expand All @@ -253,11 +267,13 @@ private Flux<String> processStreamResponse(

// Use Flux.using to manage resource lifecycle
return Flux.using(
() ->
new BufferedReader(
new InputStreamReader(inputStream, StandardCharsets.UTF_8)),
reader -> isNdjson ? readNdJsonLines(reader) : readSseLines(reader),
this::closeQuietly);
() ->
new BufferedReader(
new InputStreamReader(inputStream, StandardCharsets.UTF_8)),
reader -> isNdjson ? readNdJsonLines(reader) : readSseLines(reader),
this::closeQuietly)
// reader.lines() uses blocking I/O internally
.subscribeOn(Schedulers.boundedElastic());
}

private Flux<String> readSseLines(BufferedReader reader) {
Expand Down Expand Up @@ -310,16 +326,19 @@ public boolean isClosed() {
return closed.get();
}

private java.net.http.HttpRequest buildJdkRequest(HttpRequest request) {
private java.net.http.HttpRequest buildJdkRequest(HttpRequest request, boolean isStreaming) {
URI uri;
try {
uri = URI.create(request.getUrl());
} catch (IllegalArgumentException e) {
throw new HttpTransportException("Invalid URL: " + request.getUrl(), e);
}

var builder =
java.net.http.HttpRequest.newBuilder().uri(uri).timeout(config.getReadTimeout());
var builder = java.net.http.HttpRequest.newBuilder().uri(uri);

if (!isStreaming && config.getReadTimeout() != null) {
builder.timeout(config.getReadTimeout());
}

for (Map.Entry<String, String> header : request.getHeaders().entrySet()) {
builder.header(header.getKey(), header.getValue());
Expand Down
Loading
Loading