From c853e9afccca035ae6e2b775e13b7709a966fa0e Mon Sep 17 00:00:00 2001 From: jie65535 Date: Wed, 5 Aug 2026 00:02:29 +0800 Subject: [PATCH] llm: harden streaming responses --- src/main/kotlin/llm/ModelService.kt | 252 ++++++++++++++++-- src/main/kotlin/profile/ProfileModelClient.kt | 83 +++--- src/test/kotlin/llm/ModelServiceTest.kt | 196 ++++++++++++++ 3 files changed, 465 insertions(+), 66 deletions(-) diff --git a/src/main/kotlin/llm/ModelService.kt b/src/main/kotlin/llm/ModelService.kt index 704b616..2b8cf2f 100644 --- a/src/main/kotlin/llm/ModelService.kt +++ b/src/main/kotlin/llm/ModelService.kt @@ -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 { + 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 { - 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 = 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 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(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) diff --git a/src/main/kotlin/profile/ProfileModelClient.kt b/src/main/kotlin/profile/ProfileModelClient.kt index 3dc1cda..6047ea9 100644 --- a/src/main/kotlin/profile/ProfileModelClient.kt +++ b/src/main/kotlin/profile/ProfileModelClient.kt @@ -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, ): 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, ): 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("[\\s\\S]*?") } + + private data class CompletedProfileResponse( + val content: String, + val usage: ProfileTokenUsage, + ) } diff --git a/src/test/kotlin/llm/ModelServiceTest.kt b/src/test/kotlin/llm/ModelServiceTest.kt index 906082b..cc8ebdd 100644 --- a/src/test/kotlin/llm/ModelServiceTest.kt +++ b/src/test/kotlin/llm/ModelServiceTest.kt @@ -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 { + service.decodeStreamChunk( + """{"error":{"message":"upstream unavailable","type":"gateway_error"}}""" + ) + } + + assertTrue(failure.message.orEmpty().contains("upstream unavailable")) + } + + @Test + fun classifiesPolicyRejectionsAsNonRetryableSafetyFailures() { + val failure = assertFailsWith { + 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 { + 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 { + service.decodeStreamChunk( + """{"error":{"message":"unsupported parameter","type":"invalid_request_error"}}""" + ) + } + + assertTrue(failure !is ModelSafetyRejectionException) + } + + @Test + fun leavesAmbiguousAccountAccessErrorsRetryableForUpstreamPoolSwitching() { + val failure = assertFailsWith { + 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 { + runBlocking { + withModelResponseTimeout(10.milliseconds, "测试响应") { + delay(1.seconds) + } + } + } + } + + @Test + fun preservesOuterCoroutineTimeoutCancellation() { + assertFailsWith { + runBlocking { + withTimeout(10.milliseconds) { + withModelResponseTimeout(1.seconds, "测试响应") { + delay(1.seconds) + } + } + } + } + } }