Skip to content

Commit 3f31139

Browse files
vimtorvimtor
authored
feat(console): route migrated BYOK through provider connections (anomalyco#47266)
1 parent 475b408 commit 3f31139

4 files changed

Lines changed: 100 additions & 17 deletions

File tree

‎packages/console/app/src/lib/inference-proxy.ts‎

Lines changed: 44 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,24 @@
11
import { Resource } from "@opencode-ai/console-resource"
2-
import { Database, eq } from "@opencode-ai/console-core/drizzle/index.js"
2+
import { and, Database, eq, isNull, sql } from "@opencode-ai/console-core/drizzle/index.js"
33
import { KeyTable } from "@opencode-ai/console-core/schema/key.sql.js"
4+
import { ProviderTable } from "@opencode-ai/console-core/schema/provider.sql.js"
45
import { WorkspaceTable } from "@opencode-ai/console-core/schema/workspace.sql.js"
56

67
const paths: Record<string, string | undefined> = {
7-
"GET /zen/v1/models": "/v1/models",
88
"POST /zen/v1/chat/completions": "/openai/v1/chat/completions",
99
"POST /zen/v1/responses": "/openai/v1/responses",
1010
"POST /zen/v1/messages": "/anthropic/v1/messages",
1111
}
1212

13-
export async function proxyInference(request: Request, clientIP?: string): Promise<Response | undefined> {
13+
export async function proxyInference(
14+
request: Request,
15+
generation: {
16+
provider?: "openai" | "anthropic" | "google"
17+
/** The provider's native model ID, not the public Zen alias. */
18+
model?: string
19+
body: (model?: string) => ReadableStream<Uint8Array>
20+
},
21+
): Promise<Response | undefined> {
1422
const url = new URL(request.url)
1523
const path =
1624
paths[`${request.method} ${url.pathname}`] ??
@@ -30,23 +38,52 @@ export async function proxyInference(request: Request, clientIP?: string): Promi
3038
// Routing only; the destination owns authentication and revocation after cutover.
3139
const workspace = await Database.use((tx) =>
3240
tx
33-
.select({ migratedAt: WorkspaceTable.migrated_at })
41+
.select({
42+
id: WorkspaceTable.id,
43+
migratedAt: WorkspaceTable.migrated_at,
44+
provider: ProviderTable.provider,
45+
})
3446
.from(KeyTable)
3547
.innerJoin(WorkspaceTable, eq(WorkspaceTable.id, KeyTable.workspaceID))
48+
.leftJoin(
49+
ProviderTable,
50+
generation.provider
51+
? and(
52+
eq(ProviderTable.workspaceID, KeyTable.workspaceID),
53+
eq(ProviderTable.provider, generation.provider),
54+
isNull(ProviderTable.timeDeleted),
55+
sql`length(${ProviderTable.credentials}) > 0`,
56+
)
57+
: sql`false`,
58+
)
3659
.where(eq(KeyTable.key, key))
3760
.limit(1)
3861
.then((rows) => rows[0]),
3962
)
4063
if (!workspace?.migratedAt) return undefined
64+
const model = workspace.provider ? generation.model : undefined
65+
if (workspace.provider && !model) throw new Error("Legacy BYOK model mapping is unavailable")
4166

4267
const destination = new URL(Resource.ConsoleMigration.inferenceUrl)
43-
destination.pathname = `${destination.pathname.replace(/\/$/, "")}${path}`
68+
// Imported connections must use this same workspace/provider-derived ID.
69+
const target = model
70+
? `/custom/conn_${workspace.id.slice(4)}_${workspace.provider}${
71+
path.startsWith("/google/")
72+
? `/models/${encodeURIComponent(model)}${url.pathname.slice(url.pathname.lastIndexOf(":"))}`
73+
: url.pathname.slice("/zen/v1".length)
74+
}`
75+
: path
76+
destination.pathname = `${destination.pathname.replace(/\/$/, "")}${target}`
4477
destination.search = url.search
4578
destination.hash = ""
4679

47-
const forwarded = new Request(destination, request)
80+
// Model extraction has already read part of the body; forward its replay stream.
81+
const forwarded = new Request(
82+
destination,
83+
new Request(request, { method: request.method, body: generation.body(model) }),
84+
)
4885
forwarded.headers.set("authorization", `Bearer ${key}`)
49-
const ip = request.headers.get("cf-connecting-ip") ?? clientIP
86+
const ip = request.headers.get("cf-connecting-ip")
5087
if (ip) forwarded.headers.set("x-real-ip", ip)
5188
const requestID = request.headers.get("x-opencode-request-id") ?? request.headers.get("x-opencode-request")
5289
if (requestID) forwarded.headers.set("x-opencode-request-id", requestID)

‎packages/console/app/src/middleware.ts‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ import { createMiddleware } from "@solidjs/start/middleware"
22
import { LOCALE_HEADER, cookie, fromPathname, strip } from "~/lib/language"
33
import { normalizeReferralCode, referralCookie } from "~/lib/referral-invite"
44
import { sanitizeServerActionRequest } from "~/lib/server-action"
5-
import { proxyInference } from "~/lib/inference-proxy"
65

76
export default createMiddleware({
87
async onRequest(event) {
@@ -20,12 +19,5 @@ export default createMiddleware({
2019

2120
const referralCode = normalizeReferralCode(url.searchParams.get("ref"))
2221
if (referralCode) event.response.headers.append("set-cookie", referralCookie(referralCode))
23-
24-
return proxyInference(event.request, event.clientAddress).catch(() =>
25-
Response.json(
26-
{ error: { type: "api_error", message: "Inference routing is unavailable. Please retry later." } },
27-
{ status: 503, headers: { "Cache-Control": "no-store" } },
28-
),
29-
)
3022
},
3123
})

‎packages/console/app/src/routes/zen/util/handler.ts‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ import { countryFromRequest, isModelCountryRestricted } from "~/lib/request-coun
5050
import { isPeakPricing } from "./pricing"
5151
import { prepareRequestBody } from "./requestBody"
5252
import { requiresGoTrainingConsent } from "./trainingConsent"
53+
import { proxyInference } from "~/lib/inference-proxy"
5354

5455
type ZenData = Awaited<ReturnType<typeof ZenData.list>>
5556
type PreparedBody = Awaited<ReturnType<typeof prepareRequestBody>>
@@ -100,6 +101,26 @@ export async function handler(
100101
const ip = rawIp.includes(":") ? rawIp.split(":").slice(0, 4).join(":") : rawIp
101102
const rawZenApiKey = opts.parseApiKey(input.request.headers)
102103
const zenApiKey = rawZenApiKey === "public" ? undefined : rawZenApiKey
104+
const zenData = ZenData.list(opts.modelList)
105+
if (opts.modelList === "full" && model) {
106+
// Read routing metadata without running legacy model, auth, or balance checks.
107+
const configured = zenData.models[model]
108+
const entry = Array.isArray(configured)
109+
? configured.find((entry) => entry.formatFilter === opts.format)
110+
: configured
111+
const response = await proxyInference(input.request, {
112+
provider: entry?.byokProvider,
113+
model: entry?.providers.find((provider) => provider.id === entry.byokProvider)?.model,
114+
body: (providerModel) => requestBody?.stream(providerModel ?? model, false) ?? body,
115+
}).catch(() => {
116+
void (requestBody ? requestBody.cancel() : body.cancel()).catch(() => {})
117+
return Response.json(
118+
{ error: { type: "api_error", message: "Inference routing is unavailable. Please retry later." } },
119+
{ status: 503, headers: { "Cache-Control": "no-store" } },
120+
)
121+
})
122+
if (response) return response
123+
}
103124
const sessionId = input.request.headers.get("x-opencode-session") ?? ""
104125
const requestId = input.request.headers.get("x-opencode-request") ?? ""
105126
const ocClient = input.request.headers.get("x-opencode-client") ?? ""
@@ -112,7 +133,6 @@ export async function handler(
112133
user_agent: userAgent,
113134
"model.tier": opts.modelList === "full" ? "zen" : "go",
114135
})
115-
const zenData = ZenData.list(opts.modelList)
116136
const modelInfo = validateModel(zenData, model)
117137
const country = countryFromRequest(input.request)
118138
if (isModelCountryRestricted(modelInfo.id, country)) throw new RegionError(t("zen.api.error.countryNotAllowed"))

‎packages/console/app/src/routes/zen/v1/models.ts‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,25 @@ import { KeyTable } from "@opencode-ai/console-core/schema/key.sql.js"
55
import { WorkspaceTable } from "@opencode-ai/console-core/schema/workspace.sql.js"
66
import { ModelTable } from "@opencode-ai/console-core/schema/model.sql.js"
77
import { buildOptionsResponse, buildModelsResponse } from "~/routes/zen/util/modelsHandler"
8+
import { Resource } from "@opencode-ai/console-resource"
89

910
export async function OPTIONS(_input: APIEvent) {
1011
return buildOptionsResponse()
1112
}
1213

1314
export async function GET(input: APIEvent) {
15+
const apiKey = input.request.headers.get("authorization")?.split(" ")[1]
16+
if (apiKey && apiKey !== "public") {
17+
const response = await proxyModels(input, apiKey).catch(() =>
18+
Response.json(
19+
{ error: { type: "api_error", message: "Inference routing is unavailable. Please retry later." } },
20+
{ status: 503, headers: { "Cache-Control": "no-store" } },
21+
),
22+
)
23+
if (response) return response
24+
}
25+
1426
const disabledModels = await (() => {
15-
const apiKey = input.request.headers.get("authorization")?.split(" ")[1]
1627
if (!apiKey) return [] as string[]
1728

1829
return Database.use((tx) =>
@@ -34,3 +45,26 @@ export async function GET(input: APIEvent) {
3445

3546
return buildModelsResponse(models)
3647
}
48+
49+
async function proxyModels(input: APIEvent, apiKey: string) {
50+
// No legacy revocation or model-policy checks before destination authentication.
51+
const workspace = await Database.use((tx) =>
52+
tx
53+
.select({ migratedAt: WorkspaceTable.migrated_at })
54+
.from(KeyTable)
55+
.innerJoin(WorkspaceTable, eq(WorkspaceTable.id, KeyTable.workspaceID))
56+
.where(eq(KeyTable.key, apiKey))
57+
.limit(1)
58+
.then((rows) => rows[0]),
59+
)
60+
if (!workspace?.migratedAt) return undefined
61+
62+
const destination = new URL(Resource.ConsoleMigration.inferenceUrl)
63+
destination.pathname = `${destination.pathname.replace(/\/$/, "")}/v1/models`
64+
destination.search = new URL(input.request.url).search
65+
destination.hash = ""
66+
const headers = new Headers({ authorization: `Bearer ${apiKey}` })
67+
const ip = input.request.headers.get("cf-connecting-ip")
68+
if (ip) headers.set("x-real-ip", ip)
69+
return fetch(destination, { headers, signal: input.request.signal, redirect: "manual" })
70+
}

0 commit comments

Comments
 (0)