fix(chat): close credential leaks and harden the chat route

- Vertex: a client-supplied base URL only works with the client's own Vertex key - Accept only data: URLs for file parts in every message, so the server never downloads them - Output budget retry accounts for the thinking budget Bedrock/Anthropic add, and reads Volcengine, DashScope, SGLang and vLLM rejections; falls back to 16000 once - x-max-output-tokens can only lower the budget on server credentials - On server credentials only server models or AI_MODEL entries can be used - Drop tool results together with the invalid tool calls they belong to - Count quota tokens as input + output (cached tokens were counted twice) - Private-URL check for custom base URLs, end Langfuse traces on error/abort/early return - Fix repairToolCall ordering and placeholder, align edit_diagram prompt with operations - Panel Bedrock keys are read from ADMIN_AWS_*; forward the access code to EdgeOne - isMinimalDiagram only treats root cells as an empty canvas
2026-10-12 20:49:53 +08:00 · 2026-10-03 17:45:41 +09:00
parent a45e5b6796
commit 366480426d
11 changed files with 1066 additions and 211 deletions
@@ -0,0 +1,22 @@
+/**
+ * Check the x-access-code header against ACCESS_CODE_LIST.
+ * Returns a 401 response to send back when the check fails, or null when the
+ * request may continue (including when no access codes are configured).
+ */
+export function checkAccessCode(req: Request): Response | null {
+    const accessCodes =
+        process.env.ACCESS_CODE_LIST?.split(",")
+            .map((code) => code.trim())
+            .filter(Boolean) || []
+    if (accessCodes.length === 0) return null
+
+    const accessCodeHeader = req.headers.get("x-access-code")
+    if (accessCodeHeader && accessCodes.includes(accessCodeHeader)) return null
+
+    return Response.json(
+        {
+            error: "Invalid or missing access code. Please configure it in Settings.",
+        },
+        { status: 401 },
+    )
+}
@@ -10,6 +10,10 @@ import { aihubmix, createAihubmix } from "@aihubmix/ai-sdk-provider"
 import { fromNodeProviderChain } from "@aws-sdk/credential-providers"
 import { createOpenRouter } from "@openrouter/ai-sdk-provider"
 import { createOllama, ollama } from "ollama-ai-provider-v2"
+import {
+    adminProvidersToConfig,
+    loadAdminProviders,
+} from "@/lib/admin/providers"
 import { PROVIDER_INFO, type ProviderName } from "@/lib/types/model-config"

 export type { ProviderName }
@@ -824,8 +828,16 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
            // Use client-provided credentials if available, otherwise fall back to IAM/env vars
            const hasClientCredentials =
                overrides?.awsAccessKeyId && overrides?.awsSecretAccessKey
+            // Keys from the admin panel. The ADMIN_ names keep them out of the
+            // default AWS credential chain, which other clients such as the
+            // DynamoDB quota manager use with their own credentials.
+            const adminAccessKeyId = process.env.ADMIN_AWS_ACCESS_KEY_ID
+            const adminSecretAccessKey = process.env.ADMIN_AWS_SECRET_ACCESS_KEY
            const bedrockRegion =
-                overrides?.awsRegion || process.env.AWS_REGION || "us-west-2"
+                overrides?.awsRegion ||
+                process.env.ADMIN_AWS_REGION ||
+                process.env.AWS_REGION ||
+                "us-west-2"

            const bedrockProvider = hasClientCredentials
                ? createAmazonBedrock({
@@ -836,10 +848,16 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
                          sessionToken: overrides.awsSessionToken,
                      }),
                  })
-                : createAmazonBedrock({
-                      region: bedrockRegion,
-                      credentialProvider: fromNodeProviderChain(),
-                  })
+                : adminAccessKeyId && adminSecretAccessKey
+                  ? createAmazonBedrock({
+                        region: bedrockRegion,
+                        accessKeyId: adminAccessKeyId,
+                        secretAccessKey: adminSecretAccessKey,
+                    })
+                  : createAmazonBedrock({
+                        region: bedrockRegion,
+                        credentialProvider: fromNodeProviderChain(),
+                    })
            model = bedrockProvider(modelId)
            // Add Anthropic beta options if using Claude models via Bedrock
            if (modelId.includes("anthropic.claude")) {
@@ -872,8 +890,9 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
                // for compatibility (most proxies don't support /responses endpoint)
                const customOpenAI = createOpenAI({ apiKey, baseURL })
                model = customOpenAI.chat(modelId)
-            } else if (overrides?.apiKey) {
-                // Custom API key but official OpenAI endpoint, use Responses API
+            } else if (overrides?.apiKey || overrides?.apiKeyEnv) {
+                // Custom API key (the client's, or a server model's own env var)
+                // but official OpenAI endpoint, use Responses API
                // to support reasoning for gpt-5, o1, o3, o4 models
                const customOpenAI = createOpenAI({ apiKey })
                model = customOpenAI(modelId)
@@ -928,7 +947,9 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
                overrides?.baseUrl,
                serverBaseUrl,
            )
-            if (baseURL || overrides?.apiKey) {
+            // The default instance only reads GOOGLE_GENERATIVE_AI_API_KEY, so a
+            // server model's own env var (apiKeyEnv) needs a custom instance too
+            if (baseURL || overrides?.apiKey || overrides?.apiKeyEnv) {
                const customGoogle = createGoogleGenerativeAI({
                    apiKey,
                    ...(baseURL && { baseURL }),
@@ -941,8 +962,11 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
        }
        case "vertexai": {
            // Express Mode: Use API key for authentication
-            const vertexApiKey =
-                overrides?.vertexApiKey || process.env.GOOGLE_VERTEX_API_KEY
+            // SECURITY: a client base URL only ever gets the client's key, so the
+            // server's GOOGLE_VERTEX_API_KEY is never sent to a client-chosen host
+            const vertexApiKey = overrides?.baseUrl
+                ? overrides.vertexApiKey
+                : overrides?.vertexApiKey || process.env.GOOGLE_VERTEX_API_KEY

            if (!vertexApiKey) {
                throw new Error(
@@ -951,9 +975,13 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
                )
            }

-            // Support custom base URL from env or client override
-            const baseURL =
-                overrides?.baseUrl || process.env.GOOGLE_VERTEX_BASE_URL
+            // Support custom base URL from env or client override.
+            // A client key only goes to the client's URL or the official one.
+            const baseURL = resolveBaseURL(
+                overrides?.vertexApiKey,
+                overrides?.baseUrl,
+                process.env.GOOGLE_VERTEX_BASE_URL,
+            )

            const vertexProvider = createVertex({
                apiKey: vertexApiKey,
@@ -1079,7 +1107,7 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
                overrides?.baseUrl,
                serverBaseUrl,
            )
-            if (baseURL || overrides?.apiKey) {
+            if (baseURL || overrides?.apiKey || overrides?.apiKeyEnv) {
                const customDeepSeek = createDeepSeek({
                    apiKey,
                    ...(baseURL && { baseURL }),
@@ -1241,7 +1269,7 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
            )
            // Only use custom configuration if explicitly set (local dev or custom Gateway)
            // Otherwise undefined → AI SDK uses Vercel default (https://ai-gateway.vercel.sh/v1/ai) + OIDC
-            if (baseURL || overrides?.apiKey) {
+            if (baseURL || overrides?.apiKey || overrides?.apiKeyEnv) {
                const customGateway = createGateway({
                    apiKey,
                    ...(baseURL && { baseURL }),
@@ -1430,6 +1458,36 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
    return { model, providerOptions, headers, modelId, provider }
 }

+/**
+ * Whether the call is paid for by the server's own credentials (env keys or
+ * IAM role) rather than credentials sent with the request. Mirrors which key
+ * each branch of getAIModel ends up using.
+ */
+export function usesServerCredentials(
+    provider: ProviderName,
+    overrides?: ClientOverrides,
+): boolean {
+    switch (provider) {
+        case "bedrock":
+            return !(overrides?.awsAccessKeyId && overrides?.awsSecretAccessKey)
+        case "vertexai":
+            return !overrides?.vertexApiKey
+        case "edgeone":
+            // The platform's own endpoint, no key involved
+            return false
+        case "ollama":
+            // Only a server key costs money; a keyless local server or the
+            // client's own server does not
+            return (
+                !overrides?.baseUrl &&
+                !overrides?.apiKey &&
+                !!(overrides?.apiKeyEnv || process.env.OLLAMA_API_KEY)
+            )
+        default:
+            return !overrides?.apiKey
+    }
+}
+
 /**
 * Check if a model supports prompt caching.
 * Currently only Claude models on Bedrock support prompt caching.
@@ -1464,6 +1522,17 @@ export function getValidationModel(): ReturnType<typeof getAIModel>["model"] {
        )
    }

-    const { model } = getAIModel({ modelId })
+    // A default set in the admin panel becomes AI_PROVIDER/AI_MODEL, but its key
+    // lives in an ADMIN_-prefixed env var. Point at it the way the chat route
+    // does for server models, or the standard env var is required instead.
+    const panelDefault = adminProvidersToConfig(
+        loadAdminProviders(),
+    ).providers.find((p) => p.default && p.provider === process.env.AI_PROVIDER)
+
+    const { model } = getAIModel({
+        modelId,
+        apiKeyEnv: panelDefault?.apiKeyEnv,
+        baseUrlEnv: panelDefault?.baseUrlEnv,
+    })
    return model
 }
@@ -6,25 +6,37 @@ export const MAX_FILE_SIZE = 2 * 1024 * 1024 // 2MB
 export const MAX_FILES = 5

 // Helper function to validate file parts in messages
+// Checks every message, since history is sent to the model too
 export function validateFileParts(messages: any[]): {
    valid: boolean
    error?: string
 } {
-    const lastMessage = messages[messages.length - 1]
-    const fileParts =
-        lastMessage?.parts?.filter((p: any) => p.type === "file") || []
+    for (const message of messages) {
+        const fileParts =
+            message?.parts?.filter((p: any) => p.type === "file") || []

-    if (fileParts.length > MAX_FILES) {
-        return {
-            valid: false,
-            error: `Too many files. Maximum ${MAX_FILES} allowed.`,
+        if (fileParts.length > MAX_FILES) {
+            return {
+                valid: false,
+                error: `Too many files. Maximum ${MAX_FILES} allowed.`,
+            }
        }
-    }

-    for (const filePart of fileParts) {
-        // Data URLs format: data:image/png;base64,<data>
-        // Base64 increases size by ~33%, so we check the decoded size
-        if (filePart.url?.startsWith("data:")) {
+        for (const filePart of fileParts) {
+            // The client sends files inline. Any other URL would be downloaded
+            // by the server (AI SDK does that for models without URL support).
+            if (
+                typeof filePart.url !== "string" ||
+                !filePart.url.startsWith("data:")
+            ) {
+                return {
+                    valid: false,
+                    error: "Files must be uploaded inline as data URLs.",
+                }
+            }
+
+            // Data URLs format: data:image/png;base64,<data>
+            // Base64 increases size by ~33%, so we check the decoded size
            const base64Data = filePart.url.split(",")[1]
            if (base64Data) {
                const sizeInBytes = Math.ceil((base64Data.length * 3) / 4)
@@ -42,48 +54,89 @@ export function validateFileParts(messages: any[]): {
 }

 // Helper function to check if diagram is minimal/empty
+// Empty means no mxCell besides the root cells "0" and "1". Cells drawn in
+// draw.io get random ids, so checking for id="2" is not enough.
 export function isMinimalDiagram(xml: string): boolean {
-    const stripped = xml.replace(/\s/g, "")
-    return !stripped.includes('id="2"')
+    return !/<mxCell\b[^>]*\bid="(?![01]")/.test(xml)
+}
+
+// A tool-call input providers accept: a non-empty JSON object
+function isValidToolInput(input: unknown): boolean {
+    return !!input && typeof input === "object" && Object.keys(input).length > 0
 }

 // Helper function to replace historical tool call XML with placeholders
 // This reduces token usage and forces LLM to rely on the current diagram XML (source of truth)
-// Also fixes invalid/undefined inputs from interrupted streaming
+// Tool calls with invalid inputs are left for dropInvalidToolCalls to remove
 export function replaceHistoricalToolInputs(messages: any[]): any[] {
    return messages.map((msg) => {
        if (msg.role !== "assistant" || !Array.isArray(msg.content)) {
            return msg
        }
-        const replacedContent = msg.content
-            .map((part: any) => {
-                if (part.type === "tool-call") {
-                    const toolName = part.toolName
-                    // Fix invalid/undefined inputs from interrupted streaming
-                    if (
-                        !part.input ||
-                        typeof part.input !== "object" ||
-                        Object.keys(part.input).length === 0
-                    ) {
-                        // Skip tool calls with invalid inputs entirely
-                        return null
-                    }
-                    if (
-                        toolName === "display_diagram" ||
-                        toolName === "edit_diagram"
-                    ) {
-                        return {
-                            ...part,
-                            input: {
-                                placeholder:
-                                    "[XML content replaced - see current diagram XML in system context]",
-                            },
-                        }
-                    }
+        const replacedContent = msg.content.map((part: any) => {
+            if (
+                part.type === "tool-call" &&
+                isValidToolInput(part.input) &&
+                (part.toolName === "display_diagram" ||
+                    part.toolName === "edit_diagram")
+            ) {
+                return {
+                    ...part,
+                    input: {
+                        placeholder:
+                            "[XML content replaced - see current diagram XML in system context]",
+                    },
                }
-                return part
-            })
-            .filter(Boolean) // Remove null entries (invalid tool calls)
+            }
+            return part
+        })
        return { ...msg, content: replacedContent }
    })
 }
+
+// Remove tool-calls with invalid inputs (from failed repair or interrupted streaming),
+// together with their tool-results: providers reject a result whose call is missing.
+// Messages left empty are removed too (Bedrock rejects empty content arrays).
+export function dropInvalidToolCalls(messages: any[]): any[] {
+    const droppedIds = new Set<string>()
+    return messages
+        .map((msg) => {
+            if (!Array.isArray(msg.content)) return msg
+            const content = msg.content.filter((part: any) => {
+                if (
+                    msg.role === "assistant" &&
+                    part.type === "tool-call" &&
+                    !isValidToolInput(part.input)
+                ) {
+                    console.warn(
+                        `[chat-helpers] Dropping tool-call with invalid input:`,
+                        { toolName: part.toolName, input: part.input },
+                    )
+                    droppedIds.add(part.toolCallId)
+                    return false
+                }
+                // Results always come after their call, so the id is known by now
+                return !(
+                    part.type === "tool-result" &&
+                    droppedIds.has(part.toolCallId)
+                )
+            })
+            return { ...msg, content }
+        })
+        .filter((msg) => !Array.isArray(msg.content) || msg.content.length > 0)
+}
+
+// Fix common LLM JSON mistakes in tool-call input before jsonrepair runs
+export function fixToolInputJson(input: string): string {
+    return (
+        input
+            // Inconsistent quote escaping in XML attributes inside JSON strings:
+            // y="-20\" (opening quote unescaped, closing escaped) becomes y=\"-20\".
+            // Must run before the key fix below, which would rewrite the `="`.
+            .replace(/(\w+)="([^"]*?)\\"/g, '$1=\\"$2\\"')
+            // `:=` instead of `: `
+            .replace(/:=/g, ": ")
+            // `"key"= "` instead of `"key": "`, only for JSON keys
+            .replace(/"(\w+)"\s*=\s*"/g, '"$1": "')
+    )
+}
@@ -51,8 +51,15 @@ export function setTraceOutput(output: string) {
    if (!isLangfuseEnabled()) return

    updateActiveTrace({ output })
+    endTrace()
+}
+
+// End the observe() wrapper span (AI SDK creates its own child spans with usage).
+// It uses endOnExit: false, so every request path has to end it, or the trace
+// is never exported: stream finish, stream error/abort, and early returns.
+export function endTrace() {
+    if (!isLangfuseEnabled()) return

-    // End the observe() wrapper span (AI SDK creates its own child spans with usage)
    const activeSpan = api.trace.getActiveSpan()
    if (activeSpan) {
        activeSpan.end()
@@ -22,6 +22,12 @@ export const MAX_OUTPUT_TOKENS_LIMIT = 200000
 */
 const MIN_USABLE_OUTPUT_TOKENS = 1024

+/**
+ * Retry budget when a rejection names the budget parameter but no number we can
+ * read. It is the default from before 64000, which these providers ran with.
+ */
+const FALLBACK_OUTPUT_TOKENS = 16000
+
 /** Status codes that can carry a complaint about the requested budget. */
 const BUDGET_REJECTION_STATUSES = new Set([400, 422])

@@ -29,24 +35,8 @@ function usableLimit(value: number): number | null {
    return value >= MIN_USABLE_OUTPUT_TOKENS ? value : null
 }

-/**
- * A budget this large exceeds what some models accept. Providers reject it with a
- * 400 that names the real limit, so we parse the number out and retry once
- * instead of failing the turn.
- *
- * Formats seen in the wild:
- * - Bedrock: "The maximum tokens you requested exceeds the model limit of 4096."
- * - OpenRouter: "This endpoint's maximum context length is 64000 tokens. However,
- *   you requested about 64025 tokens (25 of text input, 64000 in the output)."
- *   Note this one is an input+output ceiling, so the input has to be subtracted.
- * - Anthropic: "max_tokens: 200000 > 64000, which is the maximum allowed..."
- * - OpenAI: "This model supports at most 16384 completion tokens"
- *
- * Every pattern names tokens explicitly. A generic one (an earlier draft matched
- * "lower than N") would reinterpret unrelated failures, and retrying on a bogus
- * number turns a readable error into an empty diagram.
- */
-export function parseOutputTokenLimit(error: unknown): number | null {
+/** Message and body of an error that may be about the budget, or null. */
+function rejectionText(error: unknown): string | null {
    const err = error as {
        message?: unknown
        responseBody?: unknown
@@ -66,24 +56,109 @@ export function parseOutputTokenLimit(error: unknown): number | null {
        typeof err?.responseBody === "string" ? err.responseBody : "",
    ].join(" ")

-    if (!text) return null
+    return text.trim() ? text : null
+}

+/**
+ * A budget this large exceeds what some models accept. Providers reject it with a
+ * 400 that names the real limit, so we parse the number out and retry once
+ * instead of failing the turn.
+ *
+ * Formats seen in the wild:
+ * - Bedrock: "The maximum tokens you requested exceeds the model limit of 4096."
+ * - OpenRouter: "This endpoint's maximum context length is 64000 tokens. However,
+ *   you requested about 64025 tokens (25 of text input, 64000 in the output)."
+ *   Note this one is an input+output ceiling, so the input has to be subtracted.
+ *   vLLM and SGLang send the same kind of ceiling, with the input written as
+ *   "6000 in the messages", "has 6000 input tokens" or "6000 tokens from the input".
+ * - Anthropic: "max_tokens: 200000 > 64000, which is the maximum allowed..."
+ * - OpenAI: "This model supports at most 16384 completion tokens"
+ * - Volcengine Ark: "The parameter `max_tokens` specified in the request are not
+ *   valid: integer above maximum value, expected a value <= 32768, but got 64000"
+ * - DashScope: "Range of max_tokens should be [1, 8192]"
+ *
+ * Every pattern names tokens explicitly. A generic one (an earlier draft matched
+ * "lower than N") would reinterpret unrelated failures, and retrying on a bogus
+ * number turns a readable error into an empty diagram.
+ */
+function readCeiling(text: string): number | null {
    // Combined input+output ceiling: subtract the input the provider counted,
    // plus a small margin because its estimate is approximate.
-    const context = text.match(/maximum context length is (\d+)/i)
+    const context = text.match(/maximum context length (?:is|of) (\d+)/i)
    if (context) {
-        const input = text.match(/(\d+) of text input/i)
-        return usableLimit(
-            Number(context[1]) - (input ? Number(input[1]) : 0) - 1024,
-        )
+        const input =
+            text.match(/(\d+) of text input/i) ||
+            text.match(/(\d+) in the messages/i) ||
+            text.match(/(\d+) tokens from the input/i) ||
+            text.match(/(\d+) input tokens/i)
+        return Number(context[1]) - (input ? Number(input[1]) : 0) - 1024
    }

    const output =
        text.match(/model limit of (\d+)/i) ||
        text.match(/> (\d+), which is the maximum/i) ||
-        text.match(/at most (\d+) completion tokens/i)
+        text.match(/at most (\d+) completion tokens/i) ||
+        text.match(/max_\w*tokens.*?expected a value (?:<=|\\u003c=) (\d+)/i) ||
+        text.match(/Range of max_tokens should be \[1,\s*(\d+)\]/i)

-    return output ? usableLimit(Number(output[1])) : null
+    return output ? Number(output[1]) : null
+}
+
+/** The usable output ceiling named in a rejection, or null. */
+export function parseOutputTokenLimit(error: unknown): number | null {
+    const text = rejectionText(error)
+    const ceiling = text ? readCeiling(text) : null
+    return ceiling === null ? null : usableLimit(ceiling)
+}
+
+/**
+ * Thinking budget the provider adds on top of maxOutputTokens. Bedrock and
+ * Anthropic send maxOutputTokens + budgetTokens as max_tokens, so a ceiling in
+ * their rejection covers both.
+ */
+function thinkingBudget(providerOptions: unknown): number {
+    const options = providerOptions as
+        | {
+              bedrock?: {
+                  reasoningConfig?: { type?: string; budgetTokens?: unknown }
+              }
+              anthropic?: {
+                  thinking?: { type?: string; budgetTokens?: unknown }
+              }
+          }
+        | undefined
+    const config =
+        options?.bedrock?.reasoningConfig ?? options?.anthropic?.thinking
+    return config?.type === "enabled" && typeof config.budgetTokens === "number"
+        ? config.budgetTokens
+        : 0
+}
+
+/**
+ * The budget to retry with after a rejection, or null to surface the error.
+ */
+export function retryOutputTokens(
+    error: unknown,
+    params: { maxOutputTokens?: number; providerOptions?: unknown },
+): number | null {
+    const requested = params.maxOutputTokens
+    const text = rejectionText(error)
+    if (!requested || !text) return null
+
+    const ceiling = readCeiling(text)
+    if (ceiling !== null) {
+        // The ceiling applies to what was actually sent, thinking included,
+        // so the retry has to leave room for the thinking too.
+        const thinking = thinkingBudget(params.providerOptions)
+        if (ceiling >= requested + thinking) return null
+        return usableLimit(ceiling - thinking)
+    }
+
+    // Names the budget parameter, but in a format we cannot read a number from
+    if (/max_\w*tokens/i.test(text) && requested > FALLBACK_OUTPUT_TOKENS) {
+        return FALLBACK_OUTPUT_TOKENS
+    }
+    return null
 }

 /**
@@ -103,17 +178,15 @@ export function withOutputTokenLimitFallback(
                try {
                    return await doStream()
                } catch (error) {
-                    const limit = parseOutputTokenLimit(error)
-                    const requested = params.maxOutputTokens
-
-                    if (!limit || !requested || limit >= requested) throw error
+                    const retry = retryOutputTokens(error, params)
+                    if (!retry) throw error

                    console.warn(
-                        `[maxOutputTokens] ${requested} rejected, retrying with ${limit}`,
+                        `[maxOutputTokens] ${params.maxOutputTokens} rejected, retrying with ${retry}`,
                    )
                    return await inner.doStream({
                        ...params,
-                        maxOutputTokens: limit,
+                        maxOutputTokens: retry,
                    })
                }
            },
@@ -135,11 +208,17 @@ function validBudget(value: string | null | undefined): number | null {
 * desktop app too), then server env, then the default. Both sources go through
 * the same validation, so a typo in either falls back instead of reaching the
 * provider.
+ *
+ * On the server's credentials the user setting can only lower the server value,
+ * so MAX_OUTPUT_TOKENS keeps capping what the server pays for.
 */
-export function resolveMaxOutputTokens(headerValue: string | null): number {
-    return (
-        validBudget(headerValue) ??
-        validBudget(process.env.MAX_OUTPUT_TOKENS) ??
-        DEFAULT_MAX_OUTPUT_TOKENS
-    )
+export function resolveMaxOutputTokens(
+    headerValue: string | null,
+    usesServerCredentials: boolean,
+): number {
+    const header = validBudget(headerValue)
+    const server =
+        validBudget(process.env.MAX_OUTPUT_TOKENS) ?? DEFAULT_MAX_OUTPUT_TOKENS
+    if (header === null) return server
+    return usesServerCredentials ? Math.min(header, server) : header
 }
@@ -41,7 +41,7 @@ parameters: {
 tool name: edit_diagram
 description: Edit specific parts of the EXISTING diagram. Use this when making small targeted changes like adding/removing elements, changing labels, or adjusting properties. This is more efficient than regenerating the entire diagram.
 parameters: {
-  edits: Array<{search: string, replace: string}>
+  operations: Array<{operation: "update" | "add" | "delete", cell_id: string, new_xml?: string}>
 }
 ---Tool3---
 tool name: append_diagram