fix(proxy): 修复 WorkBuddy/Trae 鉴权头、401、模型前缀 + 流式入口骨架

- WorkBuddyChatProxy/TraeChatProxy：Authorization 插值去掉误转义 - KtorLocalProxyServer：authorized() 失败返回 401 - 路由剥离 wbcn-/traecn- 等前缀后再转发，Trae config_name/model_name 不再带前缀 - Sub2ApiChatProxy 去掉 IOException 吞掉逻辑 - 新增 ProxyUpstreamStream 与三个代理的 openStreamingChat 入口（暂未接入实时转发）
2026-10-08 23:44:44 +08:00 · 2026-10-08 23:44:44 +08:00 · 1609e210bd
commit 1609e210bd
parent a1605e6993
5 changed files with 209 additions and 15 deletions
--- a/app/src/main/java/com/rainy/token/data/proxy/KtorLocalProxyServer.kt
+++ b/app/src/main/java/com/rainy/token/data/proxy/KtorLocalProxyServer.kt
@ -20,6 +20,9 @@ import io.ktor.server.routing.routing
 import java.io.IOException
 import javax.inject.Inject
 import javax.inject.Singleton
+import kotlinx.serialization.json.contentOrNull
+import kotlinx.serialization.json.put
+import kotlinx.serialization.json.buildJsonObject
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
@ -97,7 +100,10 @@ class KtorLocalProxyServer @Inject constructor(
                call.respond(mapOf("status" to "ok"))
            }
            get("/v1/models") {
-                if (!authorized(call, apiKey)) return@get
+                if (!authorized(call, apiKey)) {
+                    call.respond(HttpStatusCode.Unauthorized, errorBody("未授权"))
+                    return@get
+                }
                val result = sub2Api.forwardModels()
                if (result == null) {
                    call.respond(HttpStatusCode.BadRequest, errorBody("Sub2API 未配置或未登录，请在设置中填写 API Key"))
@ -106,19 +112,24 @@ class KtorLocalProxyServer @Inject constructor(
                }
            }
            post("/v1/chat/completions") {
-                if (!authorized(call, apiKey)) return@post
-                val body = call.receiveText()
-                if (body.length > MAX_REQUEST_BYTES) {
+                if (!authorized(call, apiKey)) {
+                    call.respond(HttpStatusCode.Unauthorized, errorBody("未授权"))
+                    return@post
+                }
+                val rawBody = call.receiveText()
+                if (rawBody.length > MAX_REQUEST_BYTES) {
                    call.respond(HttpStatusCode(413, ""), errorBody("请求体过大"))
                    return@post
                }
-                val model = extractModel(body)
+                val model = extractModel(rawBody)
                val route = router.route(model)
-                val conversationId = call.request.headers["X-Conversation-Id"] ?: extractUser(body)
+                val upstreamModel = stripModelPrefix(model)
+                val body = if (upstreamModel != model) rewriteModelBody(rawBody, upstreamModel) else rawBody
+                val conversationId = call.request.headers["X-Conversation-Id"] ?: extractUser(rawBody)
                val pooled = pool.next(route.kind, route.region, conversationId)
                val accountId = pooled?.accountId
                try {
-                    val result = when (route.kind) {
+                    val result: ProxyUpstreamResponse? = when (route.kind) {
                        ProviderKind.WORKBUDDY_CN, ProviderKind.WORKBUDDY_INTL ->
                            workBuddy.forwardChat(body, accountId, route.region)

@ -166,6 +177,37 @@ class KtorLocalProxyServer @Inject constructor(
    private fun contentTypeOf(raw: String): ContentType =
        runCatching { ContentType.parse(raw) }.getOrDefault(ContentType.Application.Json)

+    private fun extractStream(body: String): Boolean = runCatching {
+        val el = kotlinx.serialization.json.Json.parseToJsonElement(body)
+        (el as? kotlinx.serialization.json.JsonObject)
+            ?.get("stream")
+            ?.let { if (it is kotlinx.serialization.json.JsonPrimitive) it.contentOrNull?.toBooleanStrictOrNull() else null }
+    }.getOrNull() ?: false
+
+    private fun stripModelPrefix(model: String): String {
+        val m = model.trim()
+        val lower = m.lowercase()
+        val prefix = listOf(
+            "wbcn-", "workbuddy-cn", "codebuddy-",
+            "wbintl-", "workbuddy-intl", "workbuddy-",
+            "traeintl-", "trae-intl", "traecn-", "trae-cn", "trae-",
+            "sub2api-", "openai-",
+        ).firstOrNull { lower.startsWith(it) }
+        return if (prefix != null) m.substring(prefix.length).ifBlank { m } else m
+    }
+
+    private fun rewriteModelBody(body: String, newModel: String): String = try {
+        val obj = kotlinx.serialization.json.Json.parseToJsonElement(body)
+            as? kotlinx.serialization.json.JsonObject ?: return body
+        kotlinx.serialization.json.buildJsonObject {
+            obj.forEach { (key, value) ->
+                put(key, if (key == "model") kotlinx.serialization.json.JsonPrimitive(newModel) else value)
+            }
+        }.toString()
+    } catch (_: Throwable) {
+        body
+    }
+
    private fun extractModel(body: String): String = runCatching {
        val el = kotlinx.serialization.json.Json.parseToJsonElement(body)
        (el as? kotlinx.serialization.json.JsonObject)
--- a/app/src/main/java/com/rainy/token/data/proxy/ProxyUpstreamStream.kt
+++ b/app/src/main/java/com/rainy/token/data/proxy/ProxyUpstreamStream.kt
@ -0,0 +1,15 @@
+package com.rainy.token.data.proxy
+
+import java.io.InputStream
+
+/**
+ * 上游流式响应句柄：由本地网关在收到客户端 stream=true 请求时创建，
+ * 把上游 SSE/原始字节流实时转发给本地客户端。
+ * 使用方必须在 finally 中调用 [close]。
+ */
+class ProxyUpstreamStream(
+    val status: Int,
+    val contentType: String,
+    val input: InputStream,
+    val close: () -> Unit,
+)
--- a/app/src/main/java/com/rainy/token/data/proxy/Sub2ApiChatProxy.kt
+++ b/app/src/main/java/com/rainy/token/data/proxy/Sub2ApiChatProxy.kt
@ -34,6 +34,36 @@ class Sub2ApiChatProxy @Inject constructor(
        return forward(base = base, path = "/v1/chat/completions", requestBody = requestBody, accountId = accountId)
    }

+    /** 转发流式 POST /v1/chat/completions（上游 SSE 原样转发）。 */
+    suspend fun openStreamingChat(
+        requestBody: String,
+        accountId: String? = null,
+    ): ProxyUpstreamStream? = withContext(Dispatchers.IO) {
+        val credential = credentialRepository.get(ServiceType.SUB2API, accountId)
+            ?: return@withContext null
+        if (credential !is Credential.Sub2ApiCredential) return@withContext null
+        val auth = resolveAuth(credential) ?: return@withContext null
+        val base = normalizeBase(credential.baseUrl) ?: return@withContext null
+
+        val builder = Request.Builder()
+            .url(base + "/v1/chat/completions")
+            .addHeader("Authorization", auth)
+            .addHeader("Content-Type", "application/json")
+            .post(requestBody.toRequestBody("application/json".toMediaType()))
+
+        val response = okHttpClient.newCall(builder.build()).execute()
+        val input = response.body?.byteStream() ?: run {
+            response.close()
+            return@withContext null
+        }
+        ProxyUpstreamStream(
+            status = response.code,
+            contentType = response.header("Content-Type") ?: "text/event-stream",
+            input = input,
+            close = { response.close() },
+        )
+    }
+
    /** 转发 GET /v1/models。 */
    suspend fun forwardModels(accountId: String? = null): ProxyUpstreamResponse? {
        val base = resolveBase(accountId) ?: return null
@ -68,12 +98,8 @@ class Sub2ApiChatProxy @Inject constructor(
                .post(body.toRequestBody("application/json".toMediaType()))
        }

-        val response = try {
-            okHttpClient.newCall(builder.build()).execute()
-        } catch (_: Throwable) {
-            return@withContext null
-        }
-        val bytes = try { response.body?.bytes() ?: ByteArray(0) } catch (_: Throwable) { ByteArray(0) }
+        val response = okHttpClient.newCall(builder.build()).execute()
+        val bytes = response.body?.bytes() ?: ByteArray(0)
        val contentType = response.header("Content-Type") ?: "application/json"
        val status = response.code
        response.close()
--- a/app/src/main/java/com/rainy/token/data/proxy/TraeChatProxy.kt
+++ b/app/src/main/java/com/rainy/token/data/proxy/TraeChatProxy.kt
@ -63,7 +63,7 @@ class TraeChatProxy @Inject constructor(
                .url(base + "/api/agent/v3/llm_utils_chat")
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "text/event-stream")
-                .addHeader("Authorization", "Cloud-IDE-JWT \$jwt")
+                .addHeader("Authorization", "Cloud-IDE-JWT $jwt")
                .addHeader("X-Cloudide-Token", jwt)
                .addHeader("x-uid", credential.userId ?: "")
                .addHeader("x-device-id", credential.deviceId ?: credential.checkinDeviceId ?: "")
@ -90,6 +90,65 @@ class TraeChatProxy @Inject constructor(
            ProxyUpstreamResponse(status, contentType, bytes)
        }

+    suspend fun openStreamingChat(
+        requestBody: String,
+        accountId: String? = null,
+        forcedRegion: ProviderRegion? = null,
+    ): ProxyUpstreamStream? = withContext(Dispatchers.IO) {
+        val credential = credentialRepository.get(ServiceType.TRAE, accountId)
+            ?: return@withContext null
+        if (credential !is Credential.TraeCredential) return@withContext null
+        val jwt = credential.jwt.trim().takeIf { it.isNotBlank() }
+            ?: return@withContext null
+
+        val effectiveRegion = forcedRegion ?: runCatching {
+            ProviderRegion.valueOf(credential.region.uppercase())
+        }.getOrNull()
+        val base = if (effectiveRegion == ProviderRegion.INTL) {
+            "https://grow-normal.trae.ai"
+        } else {
+            "https://api.trae.cn"
+        }
+        val upstreamBody = buildUpstreamBody(requestBody, credential)
+        val requestId = UUID.randomUUID().toString()
+        val sessionId = UUID.randomUUID().toString().replace("-", "")
+
+        val builder = Request.Builder()
+            .url(base + "/api/agent/v3/llm_utils_chat")
+            .addHeader("Content-Type", "application/json")
+            .addHeader("Accept", "text/event-stream")
+            .addHeader("Authorization", "Cloud-IDE-JWT $jwt")
+            .addHeader("X-Cloudide-Token", jwt)
+            .addHeader("x-uid", credential.userId ?: "")
+            .addHeader("x-device-id", credential.deviceId ?: credential.checkinDeviceId ?: "")
+            .addHeader("x-machine-id", credential.deviceId ?: credential.checkinDeviceId ?: "")
+            .addHeader("x-request-id", requestId)
+            .addHeader("x-app-id", "trae")
+            .addHeader("x-ide-version", "1.0.0")
+            .addHeader("x-ide-version-code", "1000000")
+            .addHeader("x-ide-version-type", "stable")
+            .addHeader("x-os-version", "Android")
+            .addHeader("x-system-type", "Android")
+            .addHeader("User-Agent", "RainyToken/1.0")
+            .post(upstreamBody.toRequestBody("application/json".toMediaType()))
+
+        val response = try {
+            okHttpClient.newCall(builder.build()).execute()
+        } catch (e: java.io.IOException) {
+            throw e
+        }
+        val input = response.body?.byteStream() ?: run {
+            response.close()
+            return@withContext null
+        }
+        ProxyUpstreamStream(
+            status = response.code,
+            contentType = response.header("Content-Type") ?: "text/event-stream",
+            input = input,
+            close = { response.close() },
+        )
+    }
+
    private fun buildUpstreamBody(raw: String, credential: Credential.TraeCredential): String {
        val src = try {
            json.parseToJsonElement(raw) as? JsonObject
--- a/app/src/main/java/com/rainy/token/data/proxy/WorkBuddyChatProxy.kt
+++ b/app/src/main/java/com/rainy/token/data/proxy/WorkBuddyChatProxy.kt
@ -50,7 +50,7 @@ class WorkBuddyChatProxy @Inject constructor(
                .url(base + "/v2/chat/completions")
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "text/event-stream")
-                .addHeader("Authorization", "Bearer \$accessToken")
+                .addHeader("Authorization", "Bearer $accessToken")
                .addHeader("User-Agent", "RainyToken/1.0")
                .addHeader("X-IDE-Type", "VSCode")
                .addHeader("X-IDE-Name", "CodeBuddy")
@ -72,4 +72,56 @@ class WorkBuddyChatProxy @Inject constructor(
            response.close()
            ProxyUpstreamResponse(status, contentType, bytes)
        }
+    suspend fun openStreamingChat(
+        requestBody: String,
+        accountId: String? = null,
+        forcedRegion: ProviderRegion? = null,
+    ): ProxyUpstreamStream? = withContext(Dispatchers.IO) {
+        val credential = credentialRepository.get(ServiceType.WORKBUDDY, accountId)
+            ?: return@withContext null
+        if (credential !is Credential.WorkBuddyCredential) return@withContext null
+        val accessToken = credential.accessToken.trim().takeIf { it.isNotBlank() }
+            ?: return@withContext null
+
+        val effectiveRegion = forcedRegion ?: runCatching {
+            ProviderRegion.valueOf(credential.region.uppercase())
+        }.getOrNull()
+        val base = if (effectiveRegion == ProviderRegion.INTL) {
+            "https://www.workbuddy.ai"
+        } else {
+            "https://copilot.tencent.com"
+        }
+        val requestId = UUID.randomUUID().toString()
+        val builder = Request.Builder()
+            .url(base + "/v2/chat/completions")
+            .addHeader("Content-Type", "application/json")
+            .addHeader("Accept", "text/event-stream")
+            .addHeader("Authorization", "Bearer $accessToken")
+            .addHeader("User-Agent", "RainyToken/1.0")
+            .addHeader("X-IDE-Type", "VSCode")
+            .addHeader("X-IDE-Name", "CodeBuddy")
+            .addHeader("X-IDE-Version", "3.0.0")
+            .addHeader("X-Product", "CodeBuddy")
+            .addHeader("X-Agent-Intent", "craft")
+            .addHeader("X-Request-ID", requestId)
+            .addHeader("X-Conv-Request-ID", requestId)
+            .post(requestBody.toRequestBody("application/json".toMediaType()))
+
+        val response = try {
+            okHttpClient.newCall(builder.build()).execute()
+        } catch (e: java.io.IOException) {
+            throw e
+        }
+        val input = response.body?.byteStream() ?: run {
+            response.close()
+            return@withContext null
+        }
+        ProxyUpstreamStream(
+            status = response.code,
+            contentType = response.header("Content-Type") ?: "text/event-stream",
+            input = input,
+            close = { response.close() },
+        )
+    }
+
 }