zen: add safety identifier

5 files changed, 33 insertions, 26 deletions

diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts
index ee73b056d..2da542d1a 100644
--- a/packages/console/app/src/routes/zen/util/handler.ts
+++ b/packages/console/app/src/routes/zen/util/handler.ts
@@ -139,19 +139,16 @@ export async function handler(
       const startTimestamp = Date.now()
       const reqUrl = providerInfo.modifyUrl(providerInfo.api, isStream)
       const reqBody = JSON.stringify(
-        providerInfo.modifyBody(
-          {
-            ...createBodyConverter(opts.format, providerInfo.format)(body),
-            model: providerInfo.model,
-            ...(providerInfo.payloadModifier ?? {}),
-            ...Object.fromEntries(
-              Object.entries(providerInfo.payloadMappings ?? {})
-                .map(([k, v]) => [k, input.request.headers.get(v)])
-                .filter(([_k, v]) => !!v),
-            ),
-          },
-          authInfo?.workspaceID,
-        ),
+        providerInfo.modifyBody({
+          ...createBodyConverter(opts.format, providerInfo.format)(body),
+          model: providerInfo.model,
+          ...(providerInfo.payloadModifier ?? {}),
+          ...Object.fromEntries(
+            Object.entries(providerInfo.payloadMappings ?? {})
+              .map(([k, v]) => [k, input.request.headers.get(v)])
+              .filter(([_k, v]) => !!v),
+          ),
+        }),
       )
       logger.debug("REQUEST URL: " + reqUrl)
       logger.debug("REQUEST: " + reqBody.substring(0, 300) + "...")
@@ -470,15 +467,17 @@ export async function handler(
       ...(() => {
         const providerProps = zenData.providers[modelProvider.id]
         const format = providerProps.format
-        const providerModel = modelProvider.model
-        if (format === "anthropic") return anthropicHelper({ reqModel, providerModel })
-        if (format === "google") return googleHelper({ reqModel, providerModel })
-        if (format === "openai") return openaiHelper({ reqModel, providerModel })
-        return oaCompatHelper({
+        const opts = {
           reqModel,
-          providerModel,
+          providerModel: modelProvider.model,
           adjustCacheUsage: providerProps.adjustCacheUsage,
-        })
+          safetyIdentifier: ip,
+          workspaceID: authInfo?.workspaceID,
+        }
+        if (format === "anthropic") return anthropicHelper(opts)
+        if (format === "google") return googleHelper(opts)
+        if (format === "openai") return openaiHelper(opts)
+        return oaCompatHelper(opts)
       })(),
     }
   }
diff --git a/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts b/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts
index 6cb4b6a75..bdc12ba8b 100644
--- a/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts
+++ b/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts
@@ -21,17 +21,18 @@ type Usage = {
   }
 }
 
-export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage }) => ({
+export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage, safetyIdentifier }) => ({
   format: "oa-compat",
   modifyUrl: (providerApi: string) => providerApi + "/chat/completions",
   modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
     headers.set("authorization", `Bearer ${apiKey}`)
     headers.set("x-session-affinity", headers.get("x-opencode-session") ?? "")
   },
-  modifyBody: (body: Record<string, any>) => {
+  modifyBody: (body: Record<string, any>, workspaceID?: string) => {
     return {
       ...body,
       ...(body.stream ? { stream_options: { include_usage: true } } : {}),
+      ...(safetyIdentifier ? { safety_identifier: safetyIdentifier } : {}),
     }
   },
   createBinaryStreamDecoder: () => undefined,
diff --git a/packages/console/app/src/routes/zen/util/provider/openai.ts b/packages/console/app/src/routes/zen/util/provider/openai.ts
index e5649239e..3c5831a9a 100644
--- a/packages/console/app/src/routes/zen/util/provider/openai.ts
+++ b/packages/console/app/src/routes/zen/util/provider/openai.ts
@@ -12,13 +12,13 @@ type Usage = {
   total_tokens?: number
 }
 
-export const openaiHelper: ProviderHelper = () => ({
+export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({
   format: "openai",
   modifyUrl: (providerApi: string) => providerApi + "/responses",
   modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
     headers.set("authorization", `Bearer ${apiKey}`)
   },
-  modifyBody: (body: Record<string, any>, workspaceID?: string) => ({
+  modifyBody: (body: Record<string, any>) => ({
     ...body,
     ...(workspaceID ? { safety_identifier: workspaceID } : {}),
   }),
diff --git a/packages/console/app/src/routes/zen/util/provider/provider.ts b/packages/console/app/src/routes/zen/util/provider/provider.ts
index 64444ec9e..ffb23f54c 100644
--- a/packages/console/app/src/routes/zen/util/provider/provider.ts
+++ b/packages/console/app/src/routes/zen/util/provider/provider.ts
@@ -33,11 +33,17 @@ export type UsageInfo = {
   cacheWrite1hTokens?: number
 }
 
-export type ProviderHelper = (input: { reqModel: string; providerModel: string; adjustCacheUsage?: boolean }) => {
+export type ProviderHelper = (input: {
+  reqModel: string
+  providerModel: string
+  adjustCacheUsage?: boolean
+  safetyIdentifier?: string
+  workspaceID?: string
+}) => {
   format: ZenData.Format
   modifyUrl: (providerApi: string, isStream?: boolean) => string
   modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => void
-  modifyBody: (body: Record<string, any>, workspaceID?: string) => Record<string, any>
+  modifyBody: (body: Record<string, any>) => Record<string, any>
   createBinaryStreamDecoder: () => ((chunk: Uint8Array) => Uint8Array | undefined) | undefined
   streamSeparator: string
   createUsageParser: () => {
diff --git a/packages/console/core/src/model.ts b/packages/console/core/src/model.ts
index 6f28dd798..983432576 100644
--- a/packages/console/core/src/model.ts
+++ b/packages/console/core/src/model.ts
@@ -49,6 +49,7 @@ export namespace ZenData {
     payloadModifier: z.record(z.string(), z.any()).optional(),
     payloadMappings: z.record(z.string(), z.string()).optional(),
     adjustCacheUsage: z.boolean().optional(),
+    safetyIdentifier: z.boolean().optional(),
   })
 
   const ModelsSchema = z.object({