diff --git a/app/src/main/java/com/rainy/token/data/proxy/KtorLocalProxyServer.kt b/app/src/main/java/com/rainy/token/data/proxy/KtorLocalProxyServer.kt index aacdc13..48be89a 100644 --- a/app/src/main/java/com/rainy/token/data/proxy/KtorLocalProxyServer.kt +++ b/app/src/main/java/com/rainy/token/data/proxy/KtorLocalProxyServer.kt @@ -14,6 +14,7 @@ import io.ktor.server.plugins.contentnegotiation.ContentNegotiation import io.ktor.server.request.receiveText import io.ktor.server.response.respond import io.ktor.server.response.respondBytes +import io.ktor.server.response.respondOutputStream import io.ktor.server.routing.get import io.ktor.server.routing.post import io.ktor.server.routing.routing @@ -129,22 +130,58 @@ class KtorLocalProxyServer @Inject constructor( val pooled = pool.next(route.kind, route.region, conversationId) val accountId = pooled?.accountId try { - val result: ProxyUpstreamResponse? = when (route.kind) { - ProviderKind.WORKBUDDY_CN, ProviderKind.WORKBUDDY_INTL -> - workBuddy.forwardChat(body, accountId, route.region) + if (extractStream(body)) { + val stream: ProxyUpstreamStream? = when (route.kind) { + ProviderKind.WORKBUDDY_CN, ProviderKind.WORKBUDDY_INTL -> + workBuddy.openStreamingChat(body, accountId, route.region) - ProviderKind.TRAE_CN, ProviderKind.TRAE_INTL -> - trae.forwardChat(body, accountId, route.region) + ProviderKind.TRAE_CN, ProviderKind.TRAE_INTL -> + trae.openStreamingChat(body, accountId, route.region) - else -> sub2Api.forwardChat(body, accountId) - } - if (result == null) { - call.respond( - HttpStatusCode.BadRequest, - errorBody("${route.kind.displayName} 未配置或未登录,请先在设置中配置") - ) + else -> sub2Api.openStreamingChat(body, accountId) + } + if (stream == null) { + call.respond( + HttpStatusCode.BadRequest, + errorBody("${route.kind.displayName} 未配置或未登录,请先在设置中配置") + ) + } else { + call.respondOutputStream( + contentType = contentTypeOf(stream.contentType), + status = HttpStatusCode(stream.status, "") + ) { + try { + val buffer = ByteArray(8192) + val input = stream.input + while (true) { + val read = input.read(buffer) + if (read < 0) break + write(buffer, 0, read) + flush() + } + } finally { + stream.close() + } + } + } } else { - call.respondBytes(result.body, contentTypeOf(result.contentType), HttpStatusCode(result.status, "")) + val result: ProxyUpstreamResponse? = when (route.kind) { + ProviderKind.WORKBUDDY_CN, ProviderKind.WORKBUDDY_INTL -> + workBuddy.forwardChat(body, accountId, route.region) + + ProviderKind.TRAE_CN, ProviderKind.TRAE_INTL -> + trae.forwardChat(body, accountId, route.region) + + else -> sub2Api.forwardChat(body, accountId) + } + if (result == null) { + call.respond( + HttpStatusCode.BadRequest, + errorBody("${route.kind.displayName} 未配置或未登录,请先在设置中配置") + ) + } else { + call.respondBytes(result.body, contentTypeOf(result.contentType), HttpStatusCode(result.status, "")) + } } } catch (e: IOException) { val detail = e.message ?: "未知错误"