Use media-normalized context for output token capping

.changeset/fix-media-output-cap.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+Keep encoded image payload size from reducing the dynamic output token allowance.

packages/opencode/src/kilocode/session/llm.ts

@@ -6,6 +6,7 @@ import type { Provider } from "@/provider/provider"
 import { KiloSessionOverflow } from "./overflow"
 
 const SAFETY = 2048
+const MEDIA_OUTPUT_RESERVE = 2048
 const MIN_OUTPUT = 1024
 
 export namespace KiloLLM {
@@ -67,9 +68,8 @@ export namespace KiloLLM {
     const { context } = input.model.limit
     if (!context) return input.configured
 
-    const estimated =
-      input.usage?.normalized ??
-      KiloSessionOverflow.measure({ messages: input.messages, tools: input.tools }).normalized
+    const usage = input.usage ?? KiloSessionOverflow.measure({ messages: input.messages, tools: input.tools })
+    const estimated = usage.normalized + usage.media * MEDIA_OUTPUT_RESERVE
     const tokens = Math.max(input.reported ?? 0, estimated)
     const available = context - tokens - SAFETY
     // If available is ≤0 the input alone exceeds context — return the original

packages/opencode/src/kilocode/session/overflow.ts

@@ -45,10 +45,12 @@ export namespace KiloSessionOverflow {
 
   export function measure(input: Payload) {
     let extra = 0
+    let media = 0
     const normalized = JSON.stringify(input.messages, function (this: unknown, key, value: unknown) {
       if (!["data", "url", "image"].includes(key)) return value
       if (!this || typeof this !== "object" || !("type" in this)) return value
       if (!["file", "image", "media"].includes(String(this.type))) return value
+      media++
       const tokens =
         value instanceof Uint8Array
           ? Math.ceil(value.byteLength / 4)
@@ -70,6 +72,7 @@ export namespace KiloSessionOverflow {
     return {
       normalized: Math.ceil((messages + tools) * FACTOR),
       raw: Math.ceil((raw + tools) * FACTOR),
+      media,
       continuation: continued(input.messages),
     }
   }

packages/opencode/test/kilocode/session-overflow.test.ts

@@ -135,6 +135,67 @@ describe("Kilo request estimation", () => {
     expect(KiloLLM.needsEstimate({ model: mdl, configured: 32_000 })).toBe(true)
   })
 
+  test("caps output tokens from media-normalized context instead of encoded image bytes", () => {
+    const configured = 32_000
+    const messages = [
+      {
+        role: "user",
+        content: [{ type: "image", image: `data:image/png;base64,${"x".repeat(300_000)}` }],
+      },
+    ] satisfies ModelMessage[]
+    const usage = KiloSessionOverflow.measure({ messages, tools: {} })
+    const mdl = model({ context: usage.raw + 2_048 + 1_758, output: configured })
+
+    expect(usage.raw).toBeGreaterThan(usage.normalized + configured)
+    expect(KiloLLM.capOutputTokens({ model: mdl, messages, tools: {}, configured })).toBe(configured)
+  })
+
+  test("reserves output capacity for media parts without counting raw bytes", () => {
+    const configured = 32_000
+    const messages = [
+      {
+        role: "user",
+        content: [
+          { type: "image", image: `data:image/png;base64,${"x".repeat(300_000)}` },
+          { type: "image", image: `data:image/png;base64,${"y".repeat(300_000)}` },
+        ],
+      },
+    ] satisfies ModelMessage[]
+    const usage = KiloSessionOverflow.measure({ messages, tools: {} })
+    const mdl = model({ context: usage.normalized + 2_048 + 30_000, output: configured })
+
+    expect(usage.raw).toBeGreaterThan(usage.normalized + configured)
+    expect(KiloLLM.capOutputTokens({ model: mdl, messages, tools: {}, configured })).toBe(25_904)
+  })
+
+  test("reserves output capacity when usage is precomputed", () => {
+    const configured = 32_000
+    const messages = [
+      {
+        role: "user",
+        content: [
+          { type: "image", image: `data:image/png;base64,${"x".repeat(300_000)}` },
+          { type: "image", image: `data:image/png;base64,${"y".repeat(300_000)}` },
+        ],
+      },
+    ] satisfies ModelMessage[]
+    const usage = KiloSessionOverflow.measure({ messages, tools: {} })
+    const mdl = model({ context: usage.normalized + 2_048 + 30_000, output: configured })
+
+    expect(
+      KiloLLM.capOutputTokens({ model: mdl, messages, tools: {}, configured, usage }),
+    ).toBe(25_904)
+  })
+
+  test("still reduces output tokens for oversized text context", () => {
+    const configured = 32_000
+    const messages = [{ role: "user", content: "x".repeat(300_000) }] satisfies ModelMessage[]
+    const usage = KiloSessionOverflow.measure({ messages, tools: {} })
+    const mdl = model({ context: usage.normalized + 2_048 + 8_000, output: configured })
+
+    expect(KiloLLM.capOutputTokens({ model: mdl, messages, tools: {}, configured })).toBe(8_000)
+  })
+
   test("does not reduce output for encoded media payload size", () => {
     const mdl = model({ context: 200_000, output: 32_000 })
     const messages = [