llm: harden streaming responses

2026-09-15 02:56:10 +08:00 · 2026-08-05 00:02:29 +08:00
parent 1c47a69716
commit c853e9afcc
3 changed files with 465 additions and 66 deletions
@@ -17,9 +17,12 @@ import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.sync.Semaphore
 import kotlinx.coroutines.sync.withPermit
-import kotlinx.coroutines.withTimeout
+import kotlinx.coroutines.withTimeoutOrNull
 import kotlinx.serialization.json.*
 import okhttp3.Dispatcher as OkHttpDispatcher
+import okhttp3.Protocol as OkHttpProtocol
+import java.io.ByteArrayOutputStream
+import java.io.IOException
 import kotlin.time.Duration

 class ModelService(
@@ -35,15 +38,16 @@ class ModelService(

    val httpClient: HttpClient by lazy {
        HttpClient(OkHttp) {
-            this@ModelService.maxConcurrentRequests?.let { concurrencyLimit ->
            engine {
                config {
+                    protocols(MODEL_HTTP_PROTOCOLS)
+                    this@ModelService.maxConcurrentRequests?.let { concurrencyLimit ->
                        dispatcher(createRequestDispatcher(concurrencyLimit))
                    }
                }
            }
            install(HttpTimeout) {
-                // 流式响应的「首 token」与「token 间隔」超时统一由应用层 withTimeout 管控（见 chatCompletions）。
+                // 流式响应的「首 token」与「token 间隔」超时统一由应用层计时管控（见 chatCompletions）。
                // 这里特意不设 requestTimeoutMillis：否则正常但耗时较长的流式输出会被 Ktor 在中途整体掐断。
                // socket 超时作为字节级兜底，连接超时只覆盖 TCP 握手。
                socketTimeoutMillis = timeout.inWholeMilliseconds
@@ -86,24 +90,19 @@ class ModelService(
        request: ChatCompletionRequest,
        onCacheUsage: ((CacheUsage) -> Unit)? = null
    ): Flow<ChatCompletionChunk> {
-        val requestJson = json.encodeToJsonElement(ChatCompletionRequest.serializer(), request)
-            .jsonObject.toMutableMap()
-        requestJson["stream"] = JsonPrimitive(true)
-        extraBody?.forEach { (key, value) ->
-            requestJson[key] = value
-        }
-        val body = JsonObject(requestJson).toString()
+        val body = buildRequestBody(request, stream = true)

        val responseFlow: Flow<ChatCompletionChunk> = flow {
            // 关键：服务器繁忙时会拖住「响应头」，使 httpClient.post() 自身阻塞在等待响应的阶段，
-            // 因此必须把 post() 连同首个 data 块的读取一起包进 withTimeout。
+            // 因此必须把 post() 连同首个 data 块的读取一起纳入同一个应用层超时。
            // 否则首 token 超时永远不会触发（post() 还没返回，根本进不到读取循环），
            // 只能落到 Ktor 的兜底超时（很久）后再重试，表现为「等很久才报异常」。
-            // channel 在 withTimeout 外层持有：哪怕首块读取在 withTimeout 内超时，
+            // channel 在超时块外层持有：哪怕首块读取超时，
            // 只要 response.body() 已拿到通道，finally 也能释放它，避免慢速 API 重试时连接泄漏。
            var channel: ByteReadChannel? = null
+            var lineReader: LenientUtf8LineReader? = null
            try {
-                val firstDataLine = withTimeout(firstChunkTimeout) {
+                val firstDataLine = withModelResponseTimeout(firstChunkTimeout, "首个响应数据块") {
                    val response = httpClient.post("chat/completions") {
                        setBody(body)
                        contentType(ContentType.Application.Json)
@@ -115,9 +114,11 @@ class ModelService(
                    }
                    val ch: ByteReadChannel = response.body()
                    channel = ch
+                    val reader = LenientUtf8LineReader(ch)
+                    lineReader = reader
                    var found: String? = null
-                    while (currentCoroutineContext().isActive && !ch.isClosedForRead) {
-                        val line = ch.readUTF8Line() ?: continue
+                    while (currentCoroutineContext().isActive) {
+                        val line = reader.readLine() ?: break
                        if (line.startsWith("data: ")) {
                            found = line
                            break
@@ -129,19 +130,22 @@ class ModelService(

                if (firstDataLine != null && !firstDataLine.startsWith("data: [DONE]")) {
                    val firstRaw = firstDataLine.removePrefix("data: ")
-                    emit(json.decodeFromString(firstRaw))
+                    decodeStreamChunk(firstRaw)?.let { emit(it) }
                    onCacheUsage?.let { cb -> extractCacheUsage(firstRaw)?.let(cb) }

                    val ch = channel!!
-                    while (currentCoroutineContext().isActive && !ch.isClosedForRead) {
+                    val reader = lineReader!!
+                    while (currentCoroutineContext().isActive) {
                        // 流式期间同样对每次读取设「token 间隔」超时，避免中途卡死后干等兜底超时，
                        // 从而能快速失败并交给上层重试。正常流式 token 间隔远小于 firstChunkTimeout。
-                        val line = withTimeout(firstChunkTimeout) { ch.readUTF8Line() } ?: continue
+                        val line = withModelResponseTimeout(firstChunkTimeout, "流式响应数据块") {
+                            reader.readLine()
+                        } ?: break
                        when {
                            line.startsWith("data: [DONE]") -> break
                            line.startsWith("data: ") -> {
                                val raw = line.removePrefix("data: ")
-                                emit(json.decodeFromString(raw))
+                                decodeStreamChunk(raw)?.let { emit(it) }
                                onCacheUsage?.let { cb -> extractCacheUsage(raw)?.let(cb) }
                            }
                            else -> continue
@@ -154,9 +158,215 @@ class ModelService(
        }
        return responseFlow.withConcurrencyLimit(requestSemaphore)
    }
+
+    internal fun buildRequestBody(request: ChatCompletionRequest, stream: Boolean): String {
+        val requestJson = json.encodeToJsonElement(ChatCompletionRequest.serializer(), request)
+            .jsonObject.toMutableMap()
+        requestJson["stream"] = JsonPrimitive(stream)
+        extraBody?.forEach { (key, value) ->
+            requestJson[key] = value
+        }
+        return JsonObject(requestJson).toString()
    }

+    internal fun decodeStreamChunk(rawJson: String): ChatCompletionChunk? {
+        val element = try {
+            json.parseToJsonElement(rawJson)
+        } catch (cause: Exception) {
+            throw ModelStreamProtocolException(
+                "模型 SSE data 不是有效 JSON：${rawJson.toDiagnosticSnippet()}",
+                cause,
+            )
+        }
+        val payload = element as? JsonObject ?: throw ModelStreamProtocolException(
+            "模型 SSE data 必须是 JSON object：${rawJson.toDiagnosticSnippet()}"
+        )
+        if (payload.isEmpty()) return null
+
+        payload["error"]?.let { error ->
+            val errorObject = error as? JsonObject
+            val errorType = (errorObject?.get("type") as? JsonPrimitive)?.contentOrNull
+            val errorCode = (errorObject?.get("code") as? JsonPrimitive)?.contentOrNull
+            val errorMessage = (errorObject?.get("message") as? JsonPrimitive)?.contentOrNull
+            val diagnostic = "模型 SSE 返回错误事件：${error.toString().toDiagnosticSnippet()}"
+            if (isSafetyRejection(errorCode, errorMessage)) {
+                throw ModelSafetyRejectionException(errorType, errorCode, diagnostic)
+            }
+            if (errorType == "invalid_request_error" &&
+                isDeterministicRequestRejection(errorCode, errorMessage)
+            ) {
+                throw ModelRequestRejectedException(errorType, errorCode, diagnostic)
+            }
+            throw ModelStreamProtocolException(diagnostic)
+        }
+
+        val requiredChunkFields = listOf("id", "created", "model", "choices")
+        val normalized = if (requiredChunkFields.all(payload::containsKey)) {
+            payload
+        } else if ("usage" in payload) {
+            JsonObject(
+                buildMap {
+                    putAll(payload)
+                    putIfAbsent("id", JsonPrimitive("usage-only"))
+                    putIfAbsent("object", JsonPrimitive("chat.completion.chunk"))
+                    putIfAbsent("created", JsonPrimitive(0))
+                    putIfAbsent("model", JsonPrimitive("unknown"))
+                    putIfAbsent("choices", JsonArray(emptyList()))
+                }
+            )
+        } else {
+            throw ModelStreamProtocolException(
+                "模型 SSE 事件缺少 chunk 字段，keys=${payload.keys.sorted()}：" +
+                        rawJson.toDiagnosticSnippet()
+            )
+        }
+
+        return try {
+            json.decodeFromJsonElement(ChatCompletionChunk.serializer(), normalized)
+        } catch (cause: Exception) {
+            throw ModelStreamProtocolException(
+                "模型 SSE chunk 结构无效，keys=${payload.keys.sorted()}：" +
+                        rawJson.toDiagnosticSnippet(),
+                cause,
+            )
+        }
+    }
+
+    private fun String.toDiagnosticSnippet(maxChars: Int = 500): String =
+        replace(Regex("[\\r\\n]+"), " ").let { normalized ->
+            if (normalized.length <= maxChars) normalized else normalized.take(maxChars) + "...[截断]"
+        }
+
+    private fun isSafetyRejection(code: String?, message: String?): Boolean {
+        val normalizedCode = code.orEmpty().lowercase()
+        val normalizedMessage = message.orEmpty().lowercase()
+        val explicitPolicyCode = normalizedCode == "cyber_policy" ||
+                normalizedCode == "safety_policy" ||
+                normalizedCode == "moderation_blocked" ||
+                normalizedCode == "content_filter" ||
+                normalizedCode.contains("content_policy")
+        val explicitContentSafetyMessage =
+            (normalizedMessage.contains("this content") &&
+                    normalizedMessage.contains("safety reason")) ||
+                    (normalizedMessage.contains("flagged for possible") &&
+                            normalizedMessage.contains("risk")) ||
+                    normalizedMessage.contains("biological risk") ||
+                    normalizedMessage.contains("cybersecurity risk")
+        return explicitPolicyCode || explicitContentSafetyMessage
+    }
+
+    private fun isDeterministicRequestRejection(code: String?, message: String?): Boolean {
+        val normalizedCode = code.orEmpty().lowercase()
+        val normalizedMessage = message.orEmpty().lowercase()
+        return normalizedCode.contains("invalid_prompt") ||
+                normalizedCode.contains("context_length") ||
+                normalizedCode.contains("invalid_parameter") ||
+                normalizedCode.contains("unsupported_parameter") ||
+                normalizedMessage.contains("invalid prompt") ||
+                normalizedMessage.contains("maximum context length") ||
+                normalizedMessage.contains("unsupported parameter") ||
+                normalizedMessage.contains("missing required parameter")
+    }
+}
+
+internal class LenientUtf8LineReader(
+    private val channel: ByteReadChannel,
+) {
+    private val readBuffer = ByteArray(DEFAULT_BUFFER_SIZE)
+    private val lineBuffer = ByteArrayOutputStream()
+    private var readOffset = 0
+    private var readLimit = 0
+    private var skipLeadingLineFeed = false
+
+    suspend fun readLine(): String? {
+        while (true) {
+            if (readOffset >= readLimit) {
+                val read = channel.readAvailable(readBuffer)
+                if (read < 0) {
+                    channel.closedCause?.let { throw it }
+                    return if (lineBuffer.size() == 0) null else decodeBufferedLine()
+                }
+                if (read == 0) continue
+                readOffset = 0
+                readLimit = read
+            }
+
+            if (skipLeadingLineFeed) {
+                skipLeadingLineFeed = false
+                if (readBuffer[readOffset] == LINE_FEED) {
+                    readOffset++
+                    continue
+                }
+            }
+
+            val segmentStart = readOffset
+            while (readOffset < readLimit) {
+                val current = readBuffer[readOffset]
+                if (current == CARRIAGE_RETURN || current == LINE_FEED) break
+                readOffset++
+            }
+            if (readOffset > segmentStart) {
+                lineBuffer.write(readBuffer, segmentStart, readOffset - segmentStart)
+            }
+            if (readOffset >= readLimit) continue
+
+            val delimiter = readBuffer[readOffset++]
+            if (delimiter == CARRIAGE_RETURN) skipLeadingLineFeed = true
+            return decodeBufferedLine()
+        }
+    }
+
+    private fun decodeBufferedLine(): String {
+        val bytes = lineBuffer.toByteArray()
+        lineBuffer.reset()
+        // String(byte[], UTF_8) uses replacement semantics for malformed input instead of aborting the SSE stream.
+        return String(bytes, Charsets.UTF_8)
+    }
+
+    private companion object {
+        const val DEFAULT_BUFFER_SIZE = 8 * 1024
+        const val CARRIAGE_RETURN: Byte = '\r'.code.toByte()
+        const val LINE_FEED: Byte = '\n'.code.toByte()
+    }
+}
+
+internal class ModelResponseTimeoutException(
+    stage: String,
+    timeout: Duration,
+) : IOException("等待模型${stage}超过 ${timeout.inWholeMilliseconds} ms")
+
+internal class ModelStreamProtocolException(
+    message: String,
+    cause: Throwable? = null,
+) : IOException(message, cause)
+
+internal open class ModelRequestRejectedException(
+    val errorType: String?,
+    val errorCode: String?,
+    message: String,
+) : IOException(message)
+
+internal class ModelSafetyRejectionException(
+    errorType: String?,
+    errorCode: String?,
+    message: String,
+) : ModelRequestRejectedException(errorType, errorCode, message)
+
+internal suspend fun <T> withModelResponseTimeout(
+    timeout: Duration,
+    stage: String,
+    block: suspend () -> T,
+): T {
+    val completed = withTimeoutOrNull(timeout) {
+        CompletedModelResponse(block())
+    } ?: throw ModelResponseTimeoutException(stage, timeout)
+    return completed.value
+}
+
+private data class CompletedModelResponse<T>(val value: T)
+
 internal const val MAX_MODEL_CONCURRENT_REQUESTS = 512
+internal val MODEL_HTTP_PROTOCOLS = listOf(OkHttpProtocol.HTTP_1_1)

 internal fun normalizeMaxConcurrentRequests(value: Int): Int =
    value.coerceIn(1, MAX_MODEL_CONCURRENT_REQUESTS)
@@ -56,10 +56,7 @@ class ProfileModelClient(
        profile: UserProfileSnapshot,
        batch: ProfileHistoryBatch,
    ): ProfileModelResult {
-        val content = StringBuilder()
-        var lastUsage: Usage? = null
-        var cacheUsage: ModelService.CacheUsage? = null
-        endpoint.service.chatCompletions(
+        val completion = complete(
            ChatCompletionRequest(
                model = ModelId(endpoint.model),
                temperature = endpoint.temperature,
@@ -70,21 +67,13 @@ class ProfileModelClient(
                    ChatMessage.User(ProfilePromptStore.buildUserPrompt(profile, batch)),
                ),
            )
-        ) { cacheUsage = it }.collect { chunk ->
-            chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
-            chunk.usage?.let { lastUsage = it }
-        }
-
-        val raw = content.toString().replace(THINK_REGEX, "").trim()
+        )
+        val raw = completion.content.replace(THINK_REGEX, "").trim()
        val response = parseResponse(raw)
        return ProfileModelResult(
            response = response,
            rawResponse = raw,
-            usage = ProfileTokenUsage(
-                promptTokens = lastUsage?.promptTokens ?: 0,
-                completionTokens = lastUsage?.completionTokens ?: 0,
-                cachedTokens = cacheUsage?.hitTokens ?: 0,
-            ),
+            usage = completion.usage,
        )
    }

@@ -93,10 +82,7 @@ class ProfileModelClient(
        batch: ConversationProfileBatch,
        eligibleUserIds: Set<Long>,
    ): ConversationProfileModelResult {
-        val content = StringBuilder()
-        var lastUsage: Usage? = null
-        var cacheUsage: ModelService.CacheUsage? = null
-        endpoint.service.chatCompletions(
+        val completion = complete(
            ChatCompletionRequest(
                model = ModelId(endpoint.model),
                temperature = endpoint.temperature,
@@ -109,20 +95,12 @@ class ProfileModelClient(
                    ),
                ),
            )
-        ) { cacheUsage = it }.collect { chunk ->
-            chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
-            chunk.usage?.let { lastUsage = it }
-        }
-
-        val raw = content.toString().replace(THINK_REGEX, "").trim()
+        )
+        val raw = completion.content.replace(THINK_REGEX, "").trim()
        return ConversationProfileModelResult(
            response = parseObject(raw),
            rawResponse = raw,
-            usage = ProfileTokenUsage(
-                promptTokens = lastUsage?.promptTokens ?: 0,
-                completionTokens = lastUsage?.completionTokens ?: 0,
-                cachedTokens = cacheUsage?.hitTokens ?: 0,
-            ),
+            usage = completion.usage,
        )
    }

@@ -130,10 +108,7 @@ class ProfileModelClient(
        profile: UserProfileSnapshot,
        supportStats: Map<String, ProfileItemSupportStats>,
    ): ProfileCompactionModelResult {
-        val content = StringBuilder()
-        var lastUsage: Usage? = null
-        var cacheUsage: ModelService.CacheUsage? = null
-        endpoint.service.chatCompletions(
+        val completion = complete(
            ChatCompletionRequest(
                model = ModelId(endpoint.model),
                temperature = endpoint.temperature,
@@ -144,23 +119,36 @@ class ProfileModelClient(
                    ChatMessage.User(ProfilePromptStore.buildCompactionUserPrompt(profile, supportStats)),
                ),
            )
-        ) { cacheUsage = it }.collect { chunk ->
-            chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
-            chunk.usage?.let { lastUsage = it }
-        }
-
-        val raw = content.toString().replace(THINK_REGEX, "").trim()
+        )
+        val raw = completion.content.replace(THINK_REGEX, "").trim()
        return ProfileCompactionModelResult(
            response = json.decodeFromString(extractObject(raw)),
            rawResponse = raw,
-            usage = ProfileTokenUsage(
-                promptTokens = lastUsage?.promptTokens ?: 0,
-                completionTokens = lastUsage?.completionTokens ?: 0,
-                cachedTokens = cacheUsage?.hitTokens ?: 0,
-            ),
+            usage = completion.usage,
        )
    }

+    private suspend fun complete(request: ChatCompletionRequest): CompletedProfileResponse {
+        val content = StringBuilder()
+        var lastUsage: Usage? = null
+        var cacheUsage: ModelService.CacheUsage? = null
+        endpoint.service.chatCompletions(request) { cacheUsage = it }.collect { chunk ->
+            chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
+            chunk.usage?.let { lastUsage = it }
+        }
+        require(content.isNotBlank()) { "模型流式响应没有文本内容" }
+        return CompletedProfileResponse(
+            content = content.toString(),
+            usage = lastUsage.toProfileUsage(cacheUsage),
+        )
+    }
+
+    private fun Usage?.toProfileUsage(cacheUsage: ModelService.CacheUsage?) = ProfileTokenUsage(
+        promptTokens = this?.promptTokens ?: 0,
+        completionTokens = this?.completionTokens ?: 0,
+        cachedTokens = cacheUsage?.hitTokens ?: 0,
+    )
+
    private fun parseResponse(raw: String): ProfileModelResponse {
        return parseObject(raw)
    }
@@ -186,4 +174,9 @@ class ProfileModelClient(
    companion object {
        private val THINK_REGEX = Regex("<think>[\\s\\S]*?</think>")
    }
+
+    private data class CompletedProfileResponse(
+        val content: String,
+        val usage: ProfileTokenUsage,
+    )
 }
@@ -1,18 +1,33 @@
 package top.jie65535.mirai.llm

+import com.aallam.openai.api.chat.ChatCompletionRequest
+import com.aallam.openai.api.chat.ChatMessage
+import com.aallam.openai.api.model.ModelId
+import com.sun.net.httpserver.HttpServer
+import io.ktor.utils.io.ByteReadChannel
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Semaphore
 import kotlinx.coroutines.withTimeout
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.jsonObject
+import kotlinx.serialization.json.jsonPrimitive
+import okhttp3.Protocol
+import java.net.InetSocketAddress
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertNotNull
 import kotlin.test.assertNull
+import kotlin.test.assertTrue
+import kotlin.time.Duration.Companion.milliseconds
 import kotlin.time.Duration.Companion.seconds

 class ModelServiceTest {
@@ -54,12 +69,169 @@ class ModelServiceTest {
        assertEquals(128, dispatcher.maxRequestsPerHost)
    }

+    @Test
+    fun forcesHttp11ForAllModelRequests() {
+        assertEquals(listOf(Protocol.HTTP_1_1), MODEL_HTTP_PROTOCOLS)
+    }
+
    @Test
    fun clampsConcurrencyToSafetyRange() {
        assertEquals(1, normalizeMaxConcurrentRequests(0))
        assertEquals(MAX_MODEL_CONCURRENT_REQUESTS, normalizeMaxConcurrentRequests(Int.MAX_VALUE))
    }

+    @Test
+    fun buildsExplicitStreamingRequestBody() {
+        val body = service.buildRequestBody(
+            request = ChatCompletionRequest(
+                model = ModelId("test-model"),
+                messages = listOf(ChatMessage.User("hello")),
+            ),
+            stream = true,
+        )
+
+        assertEquals(
+            true,
+            Json.parseToJsonElement(body).jsonObject.getValue("stream").jsonPrimitive.content.toBoolean(),
+        )
+    }
+
+    @Test
+    fun acceptsStandaloneUsageAndEmptyHeartbeatEvents() {
+        assertNull(service.decodeStreamChunk("{}"))
+
+        val chunk = assertNotNull(
+            service.decodeStreamChunk(
+                """{"usage":{"prompt_tokens":10,"completion_tokens":4,"total_tokens":14}}"""
+            )
+        )
+        assertTrue(chunk.choices.isEmpty())
+        assertEquals(10, chunk.usage?.promptTokens)
+        assertEquals(4, chunk.usage?.completionTokens)
+    }
+
+    @Test
+    fun reportsProviderErrorEventsInsteadOfMissingChunkFields() {
+        val failure = assertFailsWith<ModelStreamProtocolException> {
+            service.decodeStreamChunk(
+                """{"error":{"message":"upstream unavailable","type":"gateway_error"}}"""
+            )
+        }
+
+        assertTrue(failure.message.orEmpty().contains("upstream unavailable"))
+    }
+
+    @Test
+    fun classifiesPolicyRejectionsAsNonRetryableSafetyFailures() {
+        val failure = assertFailsWith<ModelSafetyRejectionException> {
+            service.decodeStreamChunk(
+                """{"error":{"code":"cyber_policy","message":"This content was flagged for possible cybersecurity risk.","type":"invalid_request_error"}}"""
+            )
+        }
+
+        assertEquals("cyber_policy", failure.errorCode)
+        assertEquals("invalid_request_error", failure.errorType)
+    }
+
+    @Test
+    fun classifiesSafetyMessagesEvenWhenGatewayLabelsThemAsUpstreamErrors() {
+        val failure = assertFailsWith<ModelSafetyRejectionException> {
+            service.decodeStreamChunk(
+                """{"error":{"message":"This content was flagged for possible biological risk.","type":"upstream_error"}}"""
+            )
+        }
+
+        assertEquals("upstream_error", failure.errorType)
+        assertNull(failure.errorCode)
+    }
+
+    @Test
+    fun classifiesOtherInvalidRequestsAsNonRetryableRejections() {
+        val failure = assertFailsWith<ModelRequestRejectedException> {
+            service.decodeStreamChunk(
+                """{"error":{"message":"unsupported parameter","type":"invalid_request_error"}}"""
+            )
+        }
+
+        assertTrue(failure !is ModelSafetyRejectionException)
+    }
+
+    @Test
+    fun leavesAmbiguousAccountAccessErrorsRetryableForUpstreamPoolSwitching() {
+        val failure = assertFailsWith<ModelStreamProtocolException> {
+            service.decodeStreamChunk(
+                """{"error":{"message":"Account access is temporarily limited","type":"invalid_request_error"}}"""
+            )
+        }
+
+        assertTrue(failure.message.orEmpty().contains("temporarily limited"))
+    }
+
+    @Test
+    fun readsMalformedUtf8SseLinesWithReplacement() = runBlocking {
+        val bytes = "data: first".toByteArray() +
+                byteArrayOf(0xC3.toByte()) +
+                "\r\n\r\ndata: 第二行\n".toByteArray()
+        val reader = LenientUtf8LineReader(ByteReadChannel(bytes))
+
+        assertEquals("data: first\uFFFD", reader.readLine())
+        assertEquals("", reader.readLine())
+        assertEquals("data: 第二行", reader.readLine())
+        assertNull(reader.readLine())
+    }
+
+    @Test
+    fun drainsPrefetchedLinesAfterUnderlyingChannelCloses() = runBlocking {
+        val channel = ByteReadChannel("data: first\ndata: second\n".toByteArray())
+        val reader = LenientUtf8LineReader(channel)
+
+        assertEquals("data: first", reader.readLine())
+        assertTrue(channel.isClosedForRead)
+        assertEquals("data: second", reader.readLine())
+        assertNull(reader.readLine())
+    }
+
+    @Test
+    fun drainsCompleteSseResponseDeliveredBeforeFirstRead() = runBlocking {
+        val server = HttpServer.create(InetSocketAddress("127.0.0.1", 0), 0)
+        val responseBody = listOf(
+            """data: {"id":"chunk-1","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}""",
+            """data: {"id":"chunk-2","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}""",
+            """data: {"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}""",
+            "data: [DONE]",
+            "",
+        ).joinToString("\n\n").toByteArray()
+        server.createContext("/v1/chat/completions") { exchange ->
+            exchange.responseHeaders.add("Content-Type", "text/event-stream")
+            exchange.sendResponseHeaders(200, responseBody.size.toLong())
+            exchange.responseBody.use { body -> body.write(responseBody) }
+        }
+        server.start()
+
+        val localService = ModelService(
+            baseUrl = "http://127.0.0.1:${server.address.port}/v1/",
+            token = "test",
+            timeout = 1.seconds,
+            firstChunkTimeout = 1.seconds,
+        )
+        try {
+            val content = StringBuilder()
+            localService.chatCompletions(
+                ChatCompletionRequest(
+                    model = ModelId("test-model"),
+                    messages = listOf(ChatMessage.User("hello")),
+                )
+            ).collect { chunk ->
+                chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
+            }
+
+            assertEquals("hello world", content.toString())
+        } finally {
+            localService.httpClient.close()
+            server.stop(0)
+        }
+    }
+
    @Test
    fun queuesFlowsBeforeStartingWork() = runBlocking {
        val active = AtomicInteger()
@@ -93,4 +265,28 @@ class ModelServiceTest {
        assertEquals(2, maximumActive.get())
        assertEquals(6, started.get())
    }
+
+    @Test
+    fun convertsOwnedResponseTimeoutToRetryableException() {
+        assertFailsWith<ModelResponseTimeoutException> {
+            runBlocking {
+                withModelResponseTimeout(10.milliseconds, "测试响应") {
+                    delay(1.seconds)
+                }
+            }
+        }
+    }
+
+    @Test
+    fun preservesOuterCoroutineTimeoutCancellation() {
+        assertFailsWith<kotlinx.coroutines.TimeoutCancellationException> {
+            runBlocking {
+                withTimeout(10.milliseconds) {
+                    withModelResponseTimeout(1.seconds, "测试响应") {
+                        delay(1.seconds)
+                    }
+                }
+            }
+        }
+    }
 }