2
0
Frank 4 сар өмнө
parent
commit
57120e69ed

+ 1 - 0
infra/console.ts

@@ -102,6 +102,7 @@ const ZEN_MODELS = [
   new sst.Secret("ZEN_MODELS2"),
   new sst.Secret("ZEN_MODELS3"),
   new sst.Secret("ZEN_MODELS4"),
+  new sst.Secret("ZEN_MODELS5"),
 ]
 const STRIPE_SECRET_KEY = new sst.Secret("STRIPE_SECRET_KEY")
 const AUTH_API_URL = new sst.Linkable("AUTH_API_URL", {

+ 5 - 2
packages/console/app/src/routes/workspace/[id]/model-section.tsx

@@ -43,9 +43,12 @@ const getModelsInfo = query(async (workspaceID: string) => {
           const pA = getPriority(idA)
           const pB = getPriority(idB)
           if (pA !== pB) return pA - pB
-          return modelA.name.localeCompare(modelB.name)
+
+          const modelAName = Array.isArray(modelA) ? modelA[0].name : modelA.name
+          const modelBName = Array.isArray(modelB) ? modelB[0].name : modelB.name
+          return modelAName.localeCompare(modelBName)
         })
-        .map(([id, model]) => ({ id, name: model.name })),
+        .map(([id, model]) => ({ id, name: Array.isArray(model) ? model[0].name : model.name })),
       disabled: await Model.listDisabled(),
     }
   }, workspaceID)

+ 10 - 5
packages/console/app/src/routes/zen/util/handler.ts

@@ -57,15 +57,17 @@ export async function handler(
     const sessionId = input.request.headers.get("x-opencode-session") ?? ""
     const requestId = input.request.headers.get("x-opencode-request") ?? ""
     const projectId = input.request.headers.get("x-opencode-project") ?? ""
+    const ocClient = input.request.headers.get("x-opencode-client") ?? ""
     logger.metric({
       is_tream: isStream,
       session: sessionId,
       request: requestId,
+      client: ocClient,
     })
     const zenData = ZenData.list()
     const modelInfo = validateModel(zenData, model)
     const dataDumper = createDataDumper(sessionId, requestId, projectId)
-    const trialLimiter = createTrialLimiter(modelInfo.trial?.limit, ip)
+    const trialLimiter = createTrialLimiter(modelInfo.trial, ip, ocClient)
     const isTrial = await trialLimiter?.isTrial()
     const rateLimiter = createRateLimiter(modelInfo.id, modelInfo.rateLimit, ip)
     await rateLimiter?.check()
@@ -286,11 +288,14 @@ export async function handler(
   }
 
   function validateModel(zenData: ZenData, reqModel: string) {
-    if (!(reqModel in zenData.models)) {
-      throw new ModelError(`Model ${reqModel} not supported`)
-    }
+    if (!(reqModel in zenData.models)) throw new ModelError(`Model ${reqModel} not supported`)
+
     const modelId = reqModel as keyof typeof zenData.models
-    const modelData = zenData.models[modelId]
+    const modelData = Array.isArray(zenData.models[modelId])
+      ? zenData.models[modelId].find((model) => opts.format === model.formatFilter)
+      : zenData.models[modelId]
+
+    if (!modelData) throw new ModelError(`Model ${reqModel} not supported for format ${opts.format}`)
 
     logger.metric({ model: modelId })
 

+ 12 - 6
packages/console/app/src/routes/zen/util/trialLimiter.ts

@@ -1,12 +1,18 @@
 import { Database, eq, sql } from "@opencode-ai/console-core/drizzle/index.js"
 import { IpTable } from "@opencode-ai/console-core/schema/ip.sql.js"
 import { UsageInfo } from "./provider/provider"
+import { ZenData } from "@opencode-ai/console-core/model.js"
 
-export function createTrialLimiter(limit: number | undefined, ip: string) {
-  if (!limit) return
+export function createTrialLimiter(trial: ZenData.Trial | undefined, ip: string, client: string) {
+  if (!trial) return
   if (!ip) return
 
-  let trial: boolean
+  const limit =
+    trial.limits.find((limit) => limit.client === client)?.limit ??
+    trial.limits.find((limit) => limit.client === undefined)?.limit
+  if (!limit) return
+
+  let _isTrial: boolean
 
   return {
     isTrial: async () => {
@@ -20,11 +26,11 @@ export function createTrialLimiter(limit: number | undefined, ip: string) {
           .then((rows) => rows[0]),
       )
 
-      trial = (data?.usage ?? 0) < limit
-      return trial
+      _isTrial = (data?.usage ?? 0) < limit
+      return _isTrial
     },
     track: async (usageInfo: UsageInfo) => {
-      if (!trial) return
+      if (!_isTrial) return
       const usage =
         usageInfo.inputTokens +
         usageInfo.outputTokens +

+ 4 - 1
packages/console/core/script/promote-models.ts

@@ -16,16 +16,19 @@ const value1 = lines.find((line) => line.startsWith("ZEN_MODELS1"))?.split("=")[
 const value2 = lines.find((line) => line.startsWith("ZEN_MODELS2"))?.split("=")[1]
 const value3 = lines.find((line) => line.startsWith("ZEN_MODELS3"))?.split("=")[1]
 const value4 = lines.find((line) => line.startsWith("ZEN_MODELS4"))?.split("=")[1]
+const value5 = lines.find((line) => line.startsWith("ZEN_MODELS5"))?.split("=")[1]
 if (!value1) throw new Error("ZEN_MODELS1 not found")
 if (!value2) throw new Error("ZEN_MODELS2 not found")
 if (!value3) throw new Error("ZEN_MODELS3 not found")
 if (!value4) throw new Error("ZEN_MODELS4 not found")
+if (!value5) throw new Error("ZEN_MODELS5 not found")
 
 // validate value
-ZenData.validate(JSON.parse(value1 + value2 + value3 + value4))
+ZenData.validate(JSON.parse(value1 + value2 + value3 + value4 + value5))
 
 // update the secret
 await $`bun sst secret set ZEN_MODELS1 ${value1} --stage ${stage}`
 await $`bun sst secret set ZEN_MODELS2 ${value2} --stage ${stage}`
 await $`bun sst secret set ZEN_MODELS3 ${value3} --stage ${stage}`
 await $`bun sst secret set ZEN_MODELS4 ${value4} --stage ${stage}`
+await $`bun sst secret set ZEN_MODELS5 ${value5} --stage ${stage}`

+ 4 - 1
packages/console/core/script/pull-models.ts

@@ -16,16 +16,19 @@ const value1 = lines.find((line) => line.startsWith("ZEN_MODELS1"))?.split("=")[
 const value2 = lines.find((line) => line.startsWith("ZEN_MODELS2"))?.split("=")[1]
 const value3 = lines.find((line) => line.startsWith("ZEN_MODELS3"))?.split("=")[1]
 const value4 = lines.find((line) => line.startsWith("ZEN_MODELS4"))?.split("=")[1]
+const value5 = lines.find((line) => line.startsWith("ZEN_MODELS5"))?.split("=")[1]
 if (!value1) throw new Error("ZEN_MODELS1 not found")
 if (!value2) throw new Error("ZEN_MODELS2 not found")
 if (!value3) throw new Error("ZEN_MODELS3 not found")
 if (!value4) throw new Error("ZEN_MODELS4 not found")
+if (!value5) throw new Error("ZEN_MODELS5 not found")
 
 // validate value
-ZenData.validate(JSON.parse(value1 + value2 + value3 + value4))
+ZenData.validate(JSON.parse(value1 + value2 + value3 + value4 + value5))
 
 // update the secret
 await $`bun sst secret set ZEN_MODELS1 ${value1}`
 await $`bun sst secret set ZEN_MODELS2 ${value2}`
 await $`bun sst secret set ZEN_MODELS3 ${value3}`
 await $`bun sst secret set ZEN_MODELS4 ${value4}`
+await $`bun sst secret set ZEN_MODELS5 ${value5}`

+ 8 - 3
packages/console/core/script/update-models.ts

@@ -14,15 +14,17 @@ const oldValue1 = lines.find((line) => line.startsWith("ZEN_MODELS1"))?.split("=
 const oldValue2 = lines.find((line) => line.startsWith("ZEN_MODELS2"))?.split("=")[1]
 const oldValue3 = lines.find((line) => line.startsWith("ZEN_MODELS3"))?.split("=")[1]
 const oldValue4 = lines.find((line) => line.startsWith("ZEN_MODELS4"))?.split("=")[1]
+const oldValue5 = lines.find((line) => line.startsWith("ZEN_MODELS5"))?.split("=")[1]
 if (!oldValue1) throw new Error("ZEN_MODELS1 not found")
 if (!oldValue2) throw new Error("ZEN_MODELS2 not found")
 if (!oldValue3) throw new Error("ZEN_MODELS3 not found")
 if (!oldValue4) throw new Error("ZEN_MODELS4 not found")
+if (!oldValue5) throw new Error("ZEN_MODELS5 not found")
 
 // store the prettified json to a temp file
 const filename = `models-${Date.now()}.json`
 const tempFile = Bun.file(path.join(os.tmpdir(), filename))
-await tempFile.write(JSON.stringify(JSON.parse(oldValue1 + oldValue2 + oldValue3 + oldValue4), null, 2))
+await tempFile.write(JSON.stringify(JSON.parse(oldValue1 + oldValue2 + oldValue3 + oldValue4 + oldValue5), null, 2))
 console.log("tempFile", tempFile.name)
 
 // open temp file in vim and read the file on close
@@ -31,12 +33,15 @@ const newValue = JSON.stringify(JSON.parse(await tempFile.text()))
 ZenData.validate(JSON.parse(newValue))
 
 // update the secret
-const chunk = Math.ceil(newValue.length / 4)
+const chunk = Math.ceil(newValue.length / 5)
 const newValue1 = newValue.slice(0, chunk)
 const newValue2 = newValue.slice(chunk, chunk * 2)
 const newValue3 = newValue.slice(chunk * 2, chunk * 3)
-const newValue4 = newValue.slice(chunk * 3)
+const newValue4 = newValue.slice(chunk * 3, chunk * 4)
+const newValue5 = newValue.slice(chunk * 4)
+
 await $`bun sst secret set ZEN_MODELS1 ${newValue1}`
 await $`bun sst secret set ZEN_MODELS2 ${newValue2}`
 await $`bun sst secret set ZEN_MODELS3 ${newValue3}`
 await $`bun sst secret set ZEN_MODELS4 ${newValue4}`
+await $`bun sst secret set ZEN_MODELS5 ${newValue5}`

+ 17 - 8
packages/console/core/src/model.ts

@@ -9,7 +9,17 @@ import { Resource } from "@opencode-ai/console-resource"
 
 export namespace ZenData {
   const FormatSchema = z.enum(["anthropic", "google", "openai", "oa-compat"])
+  const TrialSchema = z.object({
+    provider: z.string(),
+    limits: z.array(
+      z.object({
+        limit: z.number(),
+        client: z.enum(["cli", "desktop"]).optional(),
+      }),
+    ),
+  })
   export type Format = z.infer<typeof FormatSchema>
+  export type Trial = z.infer<typeof TrialSchema>
 
   const ModelCostSchema = z.object({
     input: z.number(),
@@ -26,12 +36,7 @@ export namespace ZenData {
     allowAnonymous: z.boolean().optional(),
     byokProvider: z.enum(["openai", "anthropic", "google"]).optional(),
     stickyProvider: z.boolean().optional(),
-    trial: z
-      .object({
-        limit: z.number(),
-        provider: z.string(),
-      })
-      .optional(),
+    trial: TrialSchema.optional(),
     rateLimit: z.number().optional(),
     fallbackProvider: z.string().optional(),
     providers: z.array(
@@ -53,7 +58,7 @@ export namespace ZenData {
   })
 
   const ModelsSchema = z.object({
-    models: z.record(z.string(), ModelSchema),
+    models: z.record(z.string(), z.union([ModelSchema, z.array(ModelSchema.extend({ formatFilter: FormatSchema }))])),
     providers: z.record(z.string(), ProviderSchema),
   })
 
@@ -63,7 +68,11 @@ export namespace ZenData {
 
   export const list = fn(z.void(), () => {
     const json = JSON.parse(
-      Resource.ZEN_MODELS1.value + Resource.ZEN_MODELS2.value + Resource.ZEN_MODELS3.value + Resource.ZEN_MODELS4.value,
+      Resource.ZEN_MODELS1.value +
+        Resource.ZEN_MODELS2.value +
+        Resource.ZEN_MODELS3.value +
+        Resource.ZEN_MODELS4.value +
+        Resource.ZEN_MODELS5.value,
     )
     return ModelsSchema.parse(json)
   })

+ 8 - 4
packages/console/core/sst-env.d.ts

@@ -50,10 +50,6 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
-    "Enterprise": {
-      "type": "sst.cloudflare.SolidStart"
-      "url": string
-    }
     "GITHUB_APP_ID": {
       "type": "sst.sst.Secret"
       "value": string
@@ -94,6 +90,10 @@ declare module "sst" {
       "type": "sst.sst.Linkable"
       "value": string
     }
+    "Teams": {
+      "type": "sst.cloudflare.SolidStart"
+      "url": string
+    }
     "Web": {
       "type": "sst.cloudflare.Astro"
       "url": string
@@ -114,6 +114,10 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
+    "ZEN_MODELS5": {
+      "type": "sst.sst.Secret"
+      "value": string
+    }
   }
 }
 // cloudflare 

+ 8 - 4
packages/console/function/sst-env.d.ts

@@ -50,10 +50,6 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
-    "Enterprise": {
-      "type": "sst.cloudflare.SolidStart"
-      "url": string
-    }
     "GITHUB_APP_ID": {
       "type": "sst.sst.Secret"
       "value": string
@@ -94,6 +90,10 @@ declare module "sst" {
       "type": "sst.sst.Linkable"
       "value": string
     }
+    "Teams": {
+      "type": "sst.cloudflare.SolidStart"
+      "url": string
+    }
     "Web": {
       "type": "sst.cloudflare.Astro"
       "url": string
@@ -114,6 +114,10 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
+    "ZEN_MODELS5": {
+      "type": "sst.sst.Secret"
+      "value": string
+    }
   }
 }
 // cloudflare 

+ 8 - 4
packages/console/resource/sst-env.d.ts

@@ -50,10 +50,6 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
-    "Enterprise": {
-      "type": "sst.cloudflare.SolidStart"
-      "url": string
-    }
     "GITHUB_APP_ID": {
       "type": "sst.sst.Secret"
       "value": string
@@ -94,6 +90,10 @@ declare module "sst" {
       "type": "sst.sst.Linkable"
       "value": string
     }
+    "Teams": {
+      "type": "sst.cloudflare.SolidStart"
+      "url": string
+    }
     "Web": {
       "type": "sst.cloudflare.Astro"
       "url": string
@@ -114,6 +114,10 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
+    "ZEN_MODELS5": {
+      "type": "sst.sst.Secret"
+      "value": string
+    }
   }
 }
 // cloudflare 

+ 8 - 4
packages/enterprise/sst-env.d.ts

@@ -50,10 +50,6 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
-    "Enterprise": {
-      "type": "sst.cloudflare.SolidStart"
-      "url": string
-    }
     "GITHUB_APP_ID": {
       "type": "sst.sst.Secret"
       "value": string
@@ -94,6 +90,10 @@ declare module "sst" {
       "type": "sst.sst.Linkable"
       "value": string
     }
+    "Teams": {
+      "type": "sst.cloudflare.SolidStart"
+      "url": string
+    }
     "Web": {
       "type": "sst.cloudflare.Astro"
       "url": string
@@ -114,6 +114,10 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
+    "ZEN_MODELS5": {
+      "type": "sst.sst.Secret"
+      "value": string
+    }
   }
 }
 // cloudflare 

+ 8 - 4
packages/function/sst-env.d.ts

@@ -50,10 +50,6 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
-    "Enterprise": {
-      "type": "sst.cloudflare.SolidStart"
-      "url": string
-    }
     "GITHUB_APP_ID": {
       "type": "sst.sst.Secret"
       "value": string
@@ -94,6 +90,10 @@ declare module "sst" {
       "type": "sst.sst.Linkable"
       "value": string
     }
+    "Teams": {
+      "type": "sst.cloudflare.SolidStart"
+      "url": string
+    }
     "Web": {
       "type": "sst.cloudflare.Astro"
       "url": string
@@ -114,6 +114,10 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
+    "ZEN_MODELS5": {
+      "type": "sst.sst.Secret"
+      "value": string
+    }
   }
 }
 // cloudflare 

+ 8 - 4
sst-env.d.ts

@@ -65,10 +65,6 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
-    "Enterprise": {
-      "type": "sst.cloudflare.SolidStart"
-      "url": string
-    }
     "EnterpriseStorage": {
       "name": string
       "type": "sst.cloudflare.Bucket"
@@ -120,6 +116,10 @@ declare module "sst" {
       "type": "sst.sst.Linkable"
       "value": string
     }
+    "Teams": {
+      "type": "sst.cloudflare.SolidStart"
+      "url": string
+    }
     "Web": {
       "type": "sst.cloudflare.Astro"
       "url": string
@@ -140,6 +140,10 @@ declare module "sst" {
       "type": "sst.sst.Secret"
       "value": string
     }
+    "ZEN_MODELS5": {
+      "type": "sst.sst.Secret"
+      "value": string
+    }
     "ZenData": {
       "name": string
       "type": "sst.cloudflare.Bucket"