New Response.peekTrailers() API (#8921)

swankjesse · squarejesse · web-flow · commit 61a87359f64f · 2025-07-07T14:57:05.000-04:00
* New Response.peekTrailers() API This restores behavior we had in OkHttp 4.x that I broke when I changed Response.trailers() to block until trailers were available. Closes: #8916 * apiDump --------- Co-authored-by: Jesse Wilson <jwilson@squareup.com>
diff --git a/okhttp/api/android/okhttp.api b/okhttp/api/android/okhttp.api
@@ -1114,6 +1114,7 @@ public final class okhttp3/Response : java/io/Closeable {
 	public final fun networkResponse ()Lokhttp3/Response;
 	public final fun newBuilder ()Lokhttp3/Response$Builder;
 	public final fun peekBody (J)Lokhttp3/ResponseBody;
+	public final fun peekTrailers ()Lokhttp3/Headers;
 	public final fun priorResponse ()Lokhttp3/Response;
 	public final fun protocol ()Lokhttp3/Protocol;
 	public final fun receivedResponseAtMillis ()J
@@ -1219,11 +1220,16 @@ public abstract interface class okhttp3/TrailersSource {
 	public static final field Companion Lokhttp3/TrailersSource$Companion;
 	public static final field EMPTY Lokhttp3/TrailersSource;
 	public abstract fun get ()Lokhttp3/Headers;
+	public fun peek ()Lokhttp3/Headers;
 }
 
 public final class okhttp3/TrailersSource$Companion {
 }
 
+public final class okhttp3/TrailersSource$DefaultImpls {
+	public static fun peek (Lokhttp3/TrailersSource;)Lokhttp3/Headers;
+}
+
 public abstract interface class okhttp3/WebSocket {
 	public abstract fun cancel ()V
 	public abstract fun close (ILjava/lang/String;)Z
diff --git a/okhttp/api/jvm/okhttp.api b/okhttp/api/jvm/okhttp.api
@@ -1114,6 +1114,7 @@ public final class okhttp3/Response : java/io/Closeable {
 	public final fun networkResponse ()Lokhttp3/Response;
 	public final fun newBuilder ()Lokhttp3/Response$Builder;
 	public final fun peekBody (J)Lokhttp3/ResponseBody;
+	public final fun peekTrailers ()Lokhttp3/Headers;
 	public final fun priorResponse ()Lokhttp3/Response;
 	public final fun protocol ()Lokhttp3/Protocol;
 	public final fun receivedResponseAtMillis ()J
@@ -1219,11 +1220,16 @@ public abstract interface class okhttp3/TrailersSource {
 	public static final field Companion Lokhttp3/TrailersSource$Companion;
 	public static final field EMPTY Lokhttp3/TrailersSource;
 	public abstract fun get ()Lokhttp3/Headers;
+	public fun peek ()Lokhttp3/Headers;
 }
 
 public final class okhttp3/TrailersSource$Companion {
 }
 
+public final class okhttp3/TrailersSource$DefaultImpls {
+	public static fun peek (Lokhttp3/TrailersSource;)Lokhttp3/Headers;
+}
+
 public abstract interface class okhttp3/WebSocket {
 	public abstract fun cancel ()V
 	public abstract fun close (ILjava/lang/String;)Z
diff --git a/okhttp/src/commonJvmAndroid/kotlin/okhttp3/Response.kt b/okhttp/src/commonJvmAndroid/kotlin/okhttp3/Response.kt
@@ -192,6 +192,26 @@ class Response internal constructor(
   @Throws(IOException::class)
   fun trailers(): Headers = trailersSource.get()
 
+  /**
+   * Returns the trailers after the HTTP response, if they are available to read immediately. Unlike
+   * [trailers], this doesn't block if the trailers are not immediately available, and instead
+   * returns null.
+   *
+   * This will typically return null until [ResponseBody.source] has buffered the last byte of the
+   * response body. Call `body.source().request(1024 * 1024)` to block until either that's done, or
+   * 1 MiB of response data is loaded into memory. (You could use any size here, though large values
+   * risk exhausting memory.)
+   *
+   * This returns an empty value if the trailers are available, but have no data.
+   *
+   * It is not safe to call this concurrently with code that is processing the response body.
+   *
+   * @throws IOException if the trailers cannot be loaded, such as if the network connection is
+   *     dropped.
+   */
+  @Throws(IOException::class)
+  fun peekTrailers(): Headers? = trailersSource.peek()
+
   /**
    * Peeks up to [byteCount] bytes from the response body and returns them as a new response
    * body. If fewer than [byteCount] bytes are in the response body, the full response body is
diff --git a/okhttp/src/commonJvmAndroid/kotlin/okhttp3/TrailersSource.kt b/okhttp/src/commonJvmAndroid/kotlin/okhttp3/TrailersSource.kt
@@ -27,12 +27,20 @@ import okio.IOException
  * This interface is for test and production code that creates [Response] instances without making
  * an HTTP call to a remote server.
  */
-fun interface TrailersSource {
+interface TrailersSource {
+  @Throws(IOException::class)
+  fun peek(): Headers? = null
+
   @Throws(IOException::class)
   fun get(): Headers
 
   companion object {
     @JvmField
-    val EMPTY: TrailersSource = TrailersSource { Headers.EMPTY }
+    val EMPTY: TrailersSource =
+      object : TrailersSource {
+        override fun peek() = Headers.EMPTY
+
+        override fun get() = Headers.EMPTY
+      }
   }
 }
diff --git a/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/connection/Exchange.kt b/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/connection/Exchange.kt
@@ -141,7 +141,7 @@ class Exchange(
   }
 
   @Throws(IOException::class)
-  fun trailers(): Headers = codec.trailers()
+  fun peekTrailers(): Headers? = codec.peekTrailers()
 
   @Throws(SocketException::class)
   fun newWebSocketStreams(): RealWebSocket.Streams {
diff --git a/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http/CallServerInterceptor.kt b/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http/CallServerInterceptor.kt
@@ -17,8 +17,10 @@ package okhttp3.internal.http
 
 import java.io.IOException
 import java.net.ProtocolException
+import okhttp3.Headers
 import okhttp3.Interceptor
 import okhttp3.Response
+import okhttp3.TrailersSource
 import okhttp3.internal.connection.Exchange
 import okhttp3.internal.http2.ConnectionShutdownException
 import okhttp3.internal.skipAll
@@ -134,13 +136,19 @@ class CallServerInterceptor(
           response
             .newBuilder()
             .body(responseBody)
-            .trailers {
-              val source = responseBody.source()
-              if (source.isOpen) {
-                source.skipAll()
-              }
-              exchange.trailers()
-            }.build()
+            .trailers(
+              object : TrailersSource {
+                override fun peek() = exchange.peekTrailers()
+
+                override fun get(): Headers {
+                  val source = responseBody.source()
+                  if (source.isOpen) {
+                    source.skipAll()
+                  }
+                  return peek() ?: error("null trailers after exhausting response body?!")
+                }
+              },
+            ).build()
         }
       if ("close".equals(response.request.header("Connection"), ignoreCase = true) ||
         "close".equals(response.header("Connection"), ignoreCase = true)
diff --git a/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http/ExchangeCodec.kt b/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http/ExchangeCodec.kt
@@ -66,9 +66,9 @@ interface ExchangeCodec {
   @Throws(IOException::class)
   fun openResponseBodySource(response: Response): Source
 
-  /** Returns the trailers after the HTTP response. May be empty. */
+  /** Returns the trailers after the HTTP response if they're ready. May be empty. */
   @Throws(IOException::class)
-  fun trailers(): Headers
+  fun peekTrailers(): Headers?
 
   /**
    * Cancel this stream. Resources held by this stream will be cleaned up, though not synchronously.
diff --git a/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http1/Http1ExchangeCodec.kt b/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http1/Http1ExchangeCodec.kt
@@ -147,11 +147,14 @@ class Http1ExchangeCodec(
       }
     }
 
-  override fun trailers(): Headers {
+  override fun peekTrailers(): Headers? {
     if (trailers === TRAILERS_RESPONSE_BODY_TRUNCATED) {
       throw IOException("Trailers cannot be read because the response body was truncated")
     }
-    return trailers ?: error("state: $state")
+    check(state == STATE_READING_RESPONSE_BODY || state == STATE_CLOSED) {
+      "Trailers cannot be read because the state is $state"
+    }
+    return trailers
   }
 
   override fun flushRequest() {
diff --git a/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http2/Http2ExchangeCodec.kt b/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http2/Http2ExchangeCodec.kt
@@ -115,7 +115,7 @@ class Http2ExchangeCodec(
 
   override fun openResponseBodySource(response: Response): Source = stream!!.source
 
-  override fun trailers(): Headers = stream!!.trailers()
+  override fun peekTrailers(): Headers? = stream!!.peekTrailers()
 
   override fun cancel() {
     canceled = true
diff --git a/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http2/Http2Stream.kt b/okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http2/Http2Stream.kt
@@ -168,19 +168,18 @@ class Http2Stream internal constructor(
   }
 
   /**
-   * Returns the trailers. It is only safe to call this once the source stream has been completely
-   * exhausted.
+   * Returns the trailers if they're immediately available.
    */
   @Throws(IOException::class)
-  fun trailers(): Headers {
+  fun peekTrailers(): Headers? {
     withLock {
       if (source.finished && source.receiveBuffer.exhausted() && source.readBuffer.exhausted()) {
         return source.trailers ?: Headers.EMPTY
       }
       if (errorCode != null) {
         throw errorException ?: StreamResetException(errorCode!!)
       }
-      throw IllegalStateException("too early; can't read the trailers yet")
+      return null
     }
   }
 
diff --git a/okhttp/src/jvmTest/kotlin/okhttp3/ResponseJvmTest.kt b/okhttp/src/jvmTest/kotlin/okhttp3/ResponseJvmTest.kt
@@ -40,9 +40,11 @@ class ResponseJvmTest {
   fun worksIfTrailersSet() {
     val response =
       newResponse("".toResponseBody()) {
-        trailers {
-          Headers.headersOf("a", "b")
-        }
+        trailers(
+          object : TrailersSource {
+            override fun get() = Headers.headersOf("a", "b")
+          },
+        )
       }
 
     assertThat(response.trailers()["a"]).isEqualTo("b")
diff --git a/okhttp/src/jvmTest/kotlin/okhttp3/TrailersTest.kt b/okhttp/src/jvmTest/kotlin/okhttp3/TrailersTest.kt
@@ -37,6 +37,7 @@ import mockwebserver3.junit5.StartStop
 import okhttp3.Headers.Companion.headersOf
 import okhttp3.HttpUrl.Companion.toHttpUrl
 import okhttp3.ResponseBody.Companion.toResponseBody
+import okhttp3.internal.http2.Http2Connection.Companion.OKHTTP_CLIENT_WINDOW_SIZE
 import okhttp3.testing.PlatformRule
 import okio.BufferedSource
 import okio.IOException
@@ -480,18 +481,149 @@ open class TrailersTest {
 
             override fun source(): BufferedSource = error("unexpected call")
           },
-        ).trailers { headersOf("t1", "v1") }
-        .build()
+        ).trailers(
+          object : TrailersSource {
+            override fun get(): Headers = headersOf("t1", "v1")
+          },
+        ).build()
 
     assertThat(response.trailers()).isEqualTo(headersOf("t1", "v1"))
   }
 
+  @Test
+  fun peekTrailersHttp1() {
+    peekTrailers(Protocol.HTTP_1_1)
+  }
+
+  @Test
+  fun peekTrailersHttp2() {
+    peekTrailers(Protocol.HTTP_2)
+  }
+
+  private fun peekTrailers(protocol: Protocol) {
+    val responseBody = "a".repeat(OKHTTP_CLIENT_WINDOW_SIZE)
+    enableProtocol(protocol)
+
+    server.enqueue(
+      MockResponse
+        .Builder()
+        .addHeader("h1", "v1")
+        .trailers(headersOf("t1", "v2"))
+        .body(protocol, responseBody)
+        .build(),
+    )
+
+    val call = client.newCall(Request(server.url("/")))
+    call.execute().use { response ->
+      val source = response.body.source()
+      assertThat(response.header("h1")).isEqualTo("v1")
+      assertThat(response.peekTrailers()).isNull()
+      assertThat(source.readUtf8()).isEqualTo(responseBody)
+      assertThat(response.peekTrailers()).isEqualTo(headersOf("t1", "v2"))
+      assertThat(response.peekTrailers()).isEqualTo(headersOf("t1", "v2")) // Idempotent.
+      assertThat(response.trailers()).isEqualTo(headersOf("t1", "v2"))
+    }
+  }
+
+  @Test
+  fun trailersWithServerTruncatedResponseHttp1() {
+    trailersWithServerTruncatedResponse(Protocol.HTTP_1_1)
+  }
+
+  @Test
+  fun trailersWithServerTruncatedResponseHttp2() {
+    trailersWithServerTruncatedResponse(Protocol.HTTP_2)
+  }
+
+  /**
+   * If the server closes the connection while the client is consuming the response body, attempts
+   * to peek or read the trailers should throw.
+   */
+  private fun trailersWithServerTruncatedResponse(protocol: Protocol) {
+    val responseBody = "a".repeat(OKHTTP_CLIENT_WINDOW_SIZE)
+    enableProtocol(protocol)
+
+    server.enqueue(
+      MockResponse
+        .Builder()
+        .addHeader("h1", "v1")
+        .trailers(headersOf("t1", "v2"))
+        .body(protocol, responseBody)
+        .onResponseBody(CloseSocket())
+        .build(),
+    )
+
+    val call = client.newCall(Request(server.url("/")))
+    call.execute().use { response ->
+      val source = response.body.source()
+      assertThat(response.header("h1")).isEqualTo("v1")
+      assertThat(response.peekTrailers()).isNull()
+      assertFailsWith<IOException> {
+        source.readUtf8()
+      }
+      try {
+        assertThat(response.peekTrailers()).isNull() // Okay. This is what HTTP/1 does.
+      } catch (_: IOException) {
+        // Also okay. This is what HTTP/2 does.
+      }
+      assertFailsWith<IOException> {
+        response.trailers()
+      }
+    }
+  }
+
+  @Test
+  fun trailersWithClientPrematureCloseHttp1() {
+    trailersWithClientPrematureClose(Protocol.HTTP_1_1)
+  }
+
+  @Test
+  fun trailersWithClientPrematureCloseHttp2() {
+    trailersWithClientPrematureClose(Protocol.HTTP_2)
+  }
+
+  /**
+   * If the client closes the connection while it is consuming the response body, attempts to peek
+   * or read the trailers should throw.
+   */
+  private fun trailersWithClientPrematureClose(protocol: Protocol) {
+    val halfResponseBody = "a".repeat(OKHTTP_CLIENT_WINDOW_SIZE)
+    enableProtocol(protocol)
+
+    server.enqueue(
+      MockResponse
+        .Builder()
+        .addHeader("h1", "v1")
+        .trailers(headersOf("t1", "v2"))
+        .body(protocol, halfResponseBody + halfResponseBody)
+        .build(),
+    )
+
+    val call = client.newCall(Request(server.url("/")))
+    call.execute().use { response ->
+      val source = response.body.source()
+      assertThat(response.header("h1")).isEqualTo("v1")
+      assertThat(response.peekTrailers()).isNull()
+      assertThat(source.readUtf8(halfResponseBody.length.toLong())).isEqualTo(halfResponseBody)
+      source.close()
+      assertFailsWith<IllegalStateException> {
+        source.readUtf8()
+      }
+      assertFailsWith<IOException> {
+        response.peekTrailers()
+      }
+      assertFailsWith<IOException> {
+        response.trailers()
+      }
+    }
+  }
+
   private fun MockResponse.Builder.body(
     protocol: Protocol,
     body: String,
   ) = apply {
     when (protocol) {
-      Protocol.HTTP_1_1 -> chunkedBody(body)
+      Protocol.HTTP_1_1 -> chunkedBody(body, 1024) // Force multiple chunks.
       else -> body(body)
     }
   }
diff --git a/okhttp/src/jvmTest/kotlin/okhttp3/internal/http2/Http2ConnectionTest.kt b/okhttp/src/jvmTest/kotlin/okhttp3/internal/http2/Http2ConnectionTest.kt

Original file line number	Diff line number	Diff line change
`@@ -141,7 +141,7 @@ class Exchange(`
`141`	`141`	`}`
`142`	`142`
`143`	`143`	`@Throws(IOException::class)`
`144`		`- fun trailers(): Headers = codec.trailers()`
	`144`	`+ fun peekTrailers(): Headers? = codec.peekTrailers()`
`145`	`145`
`146`	`146`	`@Throws(SocketException::class)`
`147`	`147`	`fun newWebSocketStreams(): RealWebSocket.Streams {`
Original file line number	Diff line number	Diff line change
`@@ -147,11 +147,14 @@ class Http1ExchangeCodec(`
`147`	`147`	`}`
`148`	`148`	`}`
`149`	`149`
`150`		`- override fun trailers(): Headers {`
	`150`	`+ override fun peekTrailers(): Headers? {`
`151`	`151`	`if (trailers === TRAILERS_RESPONSE_BODY_TRUNCATED) {`
`152`	`152`	`throw IOException("Trailers cannot be read because the response body was truncated")`
`153`	`153`	`}`
`154`		`- return trailers ?: error("state: $state")`
	`154`	`+ check(state == STATE_READING_RESPONSE_BODY \|\| state == STATE_CLOSED) {`
	`155`	`+ "Trailers cannot be read because the state is $state"`
	`156`	`+ }`
	`157`	`+ return trailers`
`155`	`158`	`}`
`156`	`159`
`157`	`160`	`override fun flushRequest() {`
Original file line number	Diff line number	Diff line change
`@@ -168,19 +168,18 @@ class Http2Stream internal constructor(`
`168`	`168`	`}`
`169`	`169`
`170`	`170`	`/**`
`171`		`- * Returns the trailers. It is only safe to call this once the source stream has been completely`
`172`		`- * exhausted.`
	`171`	`+ * Returns the trailers if they're immediately available.`
`173`	`172`	`*/`
`174`	`173`	`@Throws(IOException::class)`
`175`		`- fun trailers(): Headers {`
	`174`	`+ fun peekTrailers(): Headers? {`
`176`	`175`	`withLock {`
`177`	`176`	`if (source.finished && source.receiveBuffer.exhausted() && source.readBuffer.exhausted()) {`
`178`	`177`	`return source.trailers ?: Headers.EMPTY`
`179`	`178`	`}`
`180`	`179`	`if (errorCode != null) {`
`181`	`180`	`throw errorException ?: StreamResetException(errorCode!!)`
`182`	`181`	`}`
`183`		`- throw IllegalStateException("too early; can't read the trailers yet")`
	`182`	`+ return null`
`184`	`183`	`}`
`185`	`184`	`}`
`186`	`185`