Kilo Max Output Plugin

kilooutputfix.ts

import type { Plugin } from "@kilocode/plugin"

/**
 * kilo-plugin-output-limit-fix
 *
 * Root cause (kilocode#10558):
 *   OUTPUT_TOKEN_MAX = 32_000 in packages/opencode/src/provider/transform.ts
 *   caps every request at max_tokens: 32000 / thinking.budget_tokens: 31999
 *   (the max thinking variant), regardless of the model's real output ceiling.
 *   Custom providers have no limit.output, so
 *   Math.min(model.limit.output, OUTPUT_TOKEN_MAX) always binds to 32k, and
 *   thinking burns the whole budget before any content is emitted
 *   (verified on OpenRouter real limit 943,717 and Ollama 131,072 —
 *   both stop at exactly completion_tokens: 32000 mid-reasoning).
 *
 * What this plugin does (via the supported `chat.params` hook):
 *   1. Raises `maxOutputTokens` from the 32k default to the model's real
 *      ceiling (model.limit.output), an explicit override, or a fallback.
 *   2. Clamps every known numeric thinking-budget field so that
 *      budget <= maxOutputTokens - reserve, leaving room for content.
 *   3. Optionally patches provider catalogs (via the `provider` hook) so
 *      custom providers with limit.output == 0 report the real ceiling.
 *
 * Downstream KiloLLM.capOutputTokens still shrinks the cap to fit context
 * (available = context - tokens - 2048, floor 1024), so overflow still
 * triggers compaction instead of provider rejection.
 */

export type OutputLimitFixOptions = {
  /** Explicit ceiling. Wins over catalog + overrides. e.g. 128000 */
  maxOutputTokens?: number
  /** Used when model.limit.output is 0/missing (custom providers). Default 128000. */
  fallbackMaxOutputTokens?: number
  /** Tokens reserved for non-reasoning content. Default 8000. */
  thinkingReserve?: number
  /** Never let a thinking budget exceed this fraction of the ceiling. Default 0.8. */
  maxThinkingRatio?: number
  /** providerID -> ceiling, e.g. { "openrouter": 943717, "ollama": 131072 } */
  providerOverrides?: Record<string, number>
  /** substring (matched against model id + api id) -> ceiling */
  modelOverrides?: Record<string, number>
  /** Only raise, never lower an already-large cap. Default true. */
  onlyRaise?: boolean
  /** Log decisions via client.app.log. Default false (debug only logs raises/clamps). */
  debug?: boolean
}

const HARDCODED_CAP = 32_000
const DEFAULT_FALLBACK = 128_000
const DEFAULT_RESERVE = 8_000
const MIN_BUDGET = 1024

type AnyModel = {
  id?: string
  providerID?: string
  limit?: { output?: number; context?: number }
  api?: { id?: string }
}

function toPositiveInt(v: unknown): number | undefined {
  const n = typeof v === "string" ? Number(v) : (v as number)
  if (typeof n === "number" && Number.isFinite(n) && n > 0) return Math.floor(n)
  return undefined
}

function matchOverride(
  model: AnyModel,
  modelOverrides?: Record<string, number>,
): number | undefined {
  if (!modelOverrides) return undefined
  const hay = `${model?.id ?? ""} ${model?.api?.id ?? ""}`.toLowerCase()
  for (const [key, value] of Object.entries(modelOverrides)) {
    if (key && hay.includes(key.toLowerCase())) {
      const n = toPositiveInt(value)
      if (n) return n
    }
  }
  return undefined
}

function resolveCeiling(model: AnyModel, opts: OutputLimitFixOptions): number {
  const explicit = toPositiveInt(opts.maxOutputTokens)
  if (explicit) return explicit

  const providerKey = (model?.providerID ?? "").toLowerCase()
  if (opts.providerOverrides) {
    for (const [key, value] of Object.entries(opts.providerOverrides)) {
      if (key && providerKey === key.toLowerCase()) {
        const n = toPositiveInt(value)
        if (n) return n
      }
    }
  }

  const modelHit = matchOverride(model, opts.modelOverrides)
  if (modelHit) return modelHit

  const catalog = toPositiveInt(model?.limit?.output)
  if (catalog) return catalog

  return toPositiveInt(opts.fallbackMaxOutputTokens) ?? DEFAULT_FALLBACK
}

function getPath(obj: any, path: string[]): any {
  let cur = obj
  for (const k of path) {
    if (cur == null || typeof cur !== "object") return undefined
    cur = cur[k]
  }
  return cur
}

function setPath(obj: any, path: string[], value: number): boolean {
  let cur = obj
  for (let i = 0; i < path.length - 1; i++) {
    const k = path[i]
    if (cur[k] == null || typeof cur[k] !== "object") return false
    cur = cur[k]
  }
  const last = path[path.length - 1]
  if (cur[last] === value) return false
  cur[last] = value
  return true
}

/** All known numeric thinking-budget locations in chat.params options. */
const BUDGET_PATHS: string[][] = [
  ["thinking", "budgetTokens"], // Anthropic
  ["reasoningConfig", "budgetTokens"], // Bedrock
  ["reasoning", "max_tokens"], // OpenRouter budget variant
  ["thinkingConfig", "thinkingBudget"], // Google
  ["thinking", "tokenBudget"], // Cohere
  ["thinkingBudget"], // Alibaba
  ["modelParams", "thinking", "budget_tokens"], // SAP Anthropic
  ["modelParams", "thinkingConfig", "thinkingBudget"], // SAP Gemini
]

/** Nested provider-key copies (defensive: some callers pre-nest options). */
const NESTED_KEYS = ["anthropic", "bedrock", "openrouter", "google", "vertex", "openaiCompatible"]

function clampBudgets(options: Record<string, any>, budgetCap: number): string[] {
  const clamped: string[] = []
  const targets: Array<{ root: any; prefix: string }> = [{ root: options, prefix: "" }]
  for (const key of NESTED_KEYS) {
    if (options[key] && typeof options[key] === "object") {
      targets.push({ root: options[key], prefix: `${key}.` })
    }
  }
  for (const { root, prefix } of targets) {
    for (const path of BUDGET_PATHS) {
      const current = getPath(root, path)
      if (typeof current === "number" && Number.isFinite(current) && current > budgetCap) {
        const next = Math.max(MIN_BUDGET, budgetCap)
        if (setPath(root, path, next)) clamped.push(`${prefix}${path.join(".")}: ${current} -> ${next}`)
      }
    }
  }
  return clamped
}

const OutputLimitFix: Plugin = async (_ctx, rawOptions) => {
  const opts = (rawOptions ?? {}) as OutputLimitFixOptions
  const reserve = toPositiveInt(opts.thinkingReserve) ?? DEFAULT_RESERVE
  const ratio =
    typeof opts.maxThinkingRatio === "number" && opts.maxThinkingRatio > 0 && opts.maxThinkingRatio < 1
      ? opts.maxThinkingRatio
      : 0.8
  const onlyRaise = opts.onlyRaise !== false
  const debug = opts.debug === true

  return {
    // Patch catalogs so custom providers (limit.output == 0) report a real
    // ceiling. This also fixes variant budget computation upstream, which
    // derives budgetTokens from model.limit.output.
    provider: {
      id: "output-limit-fix",
      async models(provider) {
        try {
          const catalog = (provider as any)?.models ?? (provider as any)?.info?.models ?? undefined
          if (!catalog || typeof catalog !== "object") return catalog as any
          const out: Record<string, any> = Array.isArray(catalog) ? {} : { ...(catalog as object) }
          const entries: Array<[string, any]> = Array.isArray(catalog)
            ? catalog.map((m: any, i: number) => [String(m?.id ?? i), m])
            : Object.entries(catalog as Record<string, any>)
          for (const [key, m] of entries) {
            const current = toPositiveInt(m?.limit?.output)
            if (current && current > HARDCODED_CAP) continue
            const ceiling = resolveCeiling(
              { id: m?.id ?? key, providerID: (provider as any)?.id, limit: m?.limit, api: m?.api },
              opts,
            )
            if (!ceiling || (current && ceiling <= current)) continue
            out[key] = { ...m, limit: { ...(m?.limit ?? {}), output: ceiling } }
          }
          return out as any
        } catch {
          return undefined as any
        }
      },
    },

    "chat.params": async (input, output) => {
      const model = input.model as unknown as AnyModel
      const options = output.options as Record<string, any>

      // Cerebras openai-compatible path drops the cap when max_completion_tokens
      // is set (transform.ts:maxOutputTokensForRequest) — don't fight it.
      if (options?.max_completion_tokens !== undefined) return
      // gpt-5 via @ai-sdk/openai-compatible proxies rejects max_tokens; the
      // core intentionally sends undefined. Respect that.
      if (output.maxOutputTokens === undefined && (model?.api?.id ?? "").toLowerCase().includes("gpt-5")) return

      const ceiling = resolveCeiling(model, opts)
      if (!ceiling) return

      const before = output.maxOutputTokens
      let after = before

      if (typeof before !== "number" || !(before > 0)) {
        after = ceiling
      } else if (before < ceiling && onlyRaise) {
        after = ceiling
      } else if (!onlyRaise && before !== ceiling) {
        after = ceiling
      }

      // Never exceed context window when known; downstream
      // KiloLLM.capOutputTokens applies the precise context fit, this is
      // just a sanity guard against absurd overrides.
      const context = toPositiveInt(model?.limit?.context)
      if (context && after && after > context) after = context

      if (after && after !== before) output.maxOutputTokens = after

      const effective = output.maxOutputTokens
      if (typeof effective === "number" && effective > 0) {
        const byReserve = effective - reserve
        const byRatio = Math.floor(effective * ratio)
        const budgetCap = Math.max(MIN_BUDGET, Math.min(byReserve, byRatio))
        const clamped = clampBudgets(options, budgetCap)
        if ((after !== before || clamped.length > 0) && (_ctx?.client as any)?.app?.log) {
          await _ctx.client.app.log({
            body: {
              service: "output-limit-fix",
              level: debug ? "debug" : "info",
              message: `maxOutputTokens ${String(before)} -> ${String(effective)}; thinking cap ${budgetCap}${clamped.length ? `; clamped ${clamped.join("; ")}` : "; thinking budgets already fit"}`,
              extra: {
                model: `${model?.providerID ?? ""}/${model?.id ?? ""}`,
                ceiling,
                clamped,
              },
            },
          }).catch(() => {})
        }
      }
    },
  }
}

export default { id: "output-limit-fix", server: OutputLimitFix }