zen: model management helper

5 files changed, 92 insertions, 27 deletions

diff --git a/packages/console/app/src/routes/zen/handler.ts b/packages/console/app/src/routes/zen/handler.ts
index 8b9a9e55f..b0f6c0972 100644
--- a/packages/console/app/src/routes/zen/handler.ts
+++ b/packages/console/app/src/routes/zen/handler.ts
@@ -10,6 +10,7 @@ import { Resource } from "@opencode/console-resource"
 import { Billing } from "../../../../core/src/billing"
 import { Actor } from "@opencode/console-core/actor.js"
 import { WorkspaceTable } from "@opencode/console-core/schema/workspace.sql.js"
+import { ZenModel } from "@opencode/console-core/model.js"
 
 export async function handler(
   input: APIEvent,
@@ -34,32 +35,7 @@ export async function handler(
   class MonthlyLimitError extends Error {}
   class ModelError extends Error {}
 
-  const ModelCostSchema = z.object({
-    input: z.number(),
-    output: z.number(),
-    cacheRead: z.number().optional(),
-    cacheWrite5m: z.number().optional(),
-    cacheWrite1h: z.number().optional(),
-  })
-
-  const ModelSchema = z.object({
-    cost: ModelCostSchema,
-    cost200K: ModelCostSchema.optional(),
-    allowAnonymous: z.boolean().optional(),
-    providers: z.array(
-      z.object({
-        id: z.string(),
-        api: z.string(),
-        apiKey: z.string(),
-        model: z.string(),
-        weight: z.number().optional(),
-        headerMappings: z.record(z.string(), z.string()).optional(),
-        disabled: z.boolean().optional(),
-      }),
-    ),
-  })
-
-  type Model = z.infer<typeof ModelSchema>
+  type Model = z.infer<typeof ZenModel.ModelSchema>
 
   const FREE_WORKSPACES = [
     "wrk_01K46JDFR0E75SG2Q8K172KF3Y", // frank
@@ -230,7 +206,7 @@ export async function handler(
   function validateModel(reqModel: string) {
     const json = JSON.parse(Resource.ZEN_MODELS.value)
 
-    const allModels = z.record(z.string(), ModelSchema).parse(json)
+    const allModels = ZenModel.ModelsSchema.parse(json)
 
     if (!(reqModel in allModels)) {
       throw new ModelError(`Model ${reqModel} not supported`)
diff --git a/packages/console/core/package.json b/packages/console/core/package.json
index 41aa0850d..06e71750b 100644
--- a/packages/console/core/package.json
+++ b/packages/console/core/package.json
@@ -20,6 +20,9 @@
     "db": "sst shell drizzle-kit",
     "db-dev": "sst shell --stage dev -- drizzle-kit",
     "db-prod": "sst shell --stage production -- drizzle-kit",
+    "update-models": "script/update-models.ts",
+    "promote-models-to-dev": "script/promote-models.ts dev",
+    "promote-models-to-prod": "script/promote-models.ts production",
     "typecheck": "tsc --noEmit"
   },
   "devDependencies": {
diff --git a/packages/console/core/script/promote-models.ts b/packages/console/core/script/promote-models.ts
new file mode 100755
index 000000000..1a5cf2fde
--- /dev/null
+++ b/packages/console/core/script/promote-models.ts
@@ -0,0 +1,24 @@
+#!/usr/bin/env bun
+
+import { $ } from "bun"
+import path from "path"
+import { ZenModel } from "../src/model"
+
+const stage = process.argv[2]
+if (!stage) throw new Error("Stage is required")
+
+const root = path.resolve(process.cwd(), "..", "..", "..")
+
+// read the secret
+const ret = await $`bun sst secret list`.cwd(root).text()
+const value = ret
+  .split("\n")
+  .find((line) => line.startsWith("ZEN_MODELS"))
+  ?.split("=")[1]
+if (!value) throw new Error("ZEN_MODELS not found")
+
+// validate value
+ZenModel.ModelsSchema.parse(JSON.parse(value))
+
+// update the secret
+await $`bun sst secret set ZEN_MODELS ${value} --stage ${stage}`
diff --git a/packages/console/core/script/update-models.ts b/packages/console/core/script/update-models.ts
new file mode 100755
index 000000000..7740fdcf8
--- /dev/null
+++ b/packages/console/core/script/update-models.ts
@@ -0,0 +1,32 @@
+#!/usr/bin/env bun
+
+import { $ } from "bun"
+import path from "path"
+import os from "os"
+import { ZenModel } from "../src/model"
+
+const root = path.resolve(process.cwd(), "..", "..", "..")
+const models = await $`bun sst secret list`.cwd(root).text()
+console.log("models", models)
+
+// read the line starting with "ZEN_MODELS"
+const oldValue = models
+  .split("\n")
+  .find((line) => line.startsWith("ZEN_MODELS"))
+  ?.split("=")[1]
+if (!oldValue) throw new Error("ZEN_MODELS not found")
+console.log("oldValue", oldValue)
+
+// 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(oldValue), null, 2))
+console.log("tempFile", tempFile.name)
+
+// open temp file in vim and read the file on close
+await $`vim ${tempFile.name}`
+const newValue = JSON.parse(await tempFile.text())
+ZenModel.ModelsSchema.parse(newValue)
+
+// update the secret
+await $`bun sst secret set ZEN_MODELS ${JSON.stringify(newValue)}`
diff --git a/packages/console/core/src/model.ts b/packages/console/core/src/model.ts
new file mode 100644
index 000000000..028d94655
--- /dev/null
+++ b/packages/console/core/src/model.ts
@@ -0,0 +1,30 @@
+import { z } from "zod"
+
+export namespace ZenModel {
+  const ModelCostSchema = z.object({
+    input: z.number(),
+    output: z.number(),
+    cacheRead: z.number().optional(),
+    cacheWrite5m: z.number().optional(),
+    cacheWrite1h: z.number().optional(),
+  })
+
+  export const ModelSchema = z.object({
+    cost: ModelCostSchema,
+    cost200K: ModelCostSchema.optional(),
+    allowAnonymous: z.boolean().optional(),
+    providers: z.array(
+      z.object({
+        id: z.string(),
+        api: z.string(),
+        apiKey: z.string(),
+        model: z.string(),
+        weight: z.number().optional(),
+        headerMappings: z.record(z.string(), z.string()).optional(),
+        disabled: z.boolean().optional(),
+      }),
+    ),
+  })
+
+  export const ModelsSchema = z.record(z.string(), ModelSchema)
+}