Merge pull request #513 from isac322/fix/input-budget-as-percentage-base

fix: use input budget (limit.input or context − output) as percentage base

Author: Tarquinen

lib/hooks.ts

@@ -13,6 +13,7 @@ import {
     stripHallucinationsFromString,
     stripStaleMetadata,
     syncCompressionBlocks,
+    computeInputBudget,
 } from "./messages"
 import { renderSystemPrompt, type PromptStore } from "./prompts"
 import { buildProtectedToolsExtension } from "./prompts/extensions/system"
@@ -53,11 +54,17 @@ export function createSystemPromptHandler(
     prompts: PromptStore,
 ) {
     return async (
-        input: { sessionID?: string; model: { limit: { context: number } } },
+        input: {
+            sessionID?: string
+            model: { limit: { context: number; input?: number; output?: number } }
+        },
         output: { system: string[] },
     ) => {
         if (input.model?.limit?.context) {
-            state.modelContextLimit = input.model.limit.context
+            const inputBudget = computeInputBudget(input.model.limit)
+            if (inputBudget !== undefined) {
+                state.modelContextLimit = inputBudget
+            }
             logger.debug("Cached model context limit", { limit: state.modelContextLimit })
         }
 

lib/messages/index.ts

@@ -1,6 +1,7 @@
 export { prune } from "./prune"
 export { syncCompressionBlocks } from "./sync"
 export { injectCompressNudges } from "./inject/inject"
+export { computeInputBudget } from "./inject/utils"
 export { injectMessageIds } from "./inject/inject"
 export { injectExtendedSubAgentResults } from "./inject/subagent-results"
 export { stripStaleMetadata } from "./reasoning-strip"

lib/messages/inject/utils.ts

@@ -34,6 +34,20 @@ export interface LastNonIgnoredMessage {
     index: number
 }
 
+interface ModelLimit {
+    context: number
+    input?: number
+    output?: number
+}
+
+export function computeInputBudget(limit: ModelLimit): number | undefined {
+    if (!limit.context) {
+        return undefined
+    }
+
+    return limit.input ?? Math.max(0, limit.context - (limit.output ?? 0))
+}
+
 export function getNudgeFrequency(config: PluginConfig): number {
     return Math.max(1, Math.floor(config.compress.nudgeFrequency || 1))
 }

tests/input-budget.test.ts

@@ -0,0 +1,38 @@
+import assert from "node:assert/strict"
+import test from "node:test"
+import { computeInputBudget } from "../lib/messages/inject/utils"
+
+test("computeInputBudget uses limit.input when defined (split-budget OpenAI models)", () => {
+    // gpt-5.4-mini, gpt-5.5: 400K context, 272K input, 128K output
+    assert.equal(computeInputBudget({ context: 400000, input: 272000, output: 128000 }), 272000)
+    // gpt-5.4: 1.05M context, 922K input, 128K output
+    assert.equal(computeInputBudget({ context: 1050000, input: 922000, output: 128000 }), 922000)
+})
+
+test("computeInputBudget subtracts output from context when limit.input is undefined (shared-pool models)", () => {
+    // claude-opus-4-7: 1M context, 128K output, no explicit input limit
+    assert.equal(computeInputBudget({ context: 1000000, output: 128000 }), 872000)
+    // claude-haiku-4-5: 200K context, 64K output
+    assert.equal(computeInputBudget({ context: 200000, output: 64000 }), 136000)
+    // gpt-4o: 128K context, 16384 output
+    assert.equal(computeInputBudget({ context: 128000, output: 16384 }), 111616)
+})
+
+test("computeInputBudget treats missing output as 0", () => {
+    assert.equal(computeInputBudget({ context: 200000 }), 200000)
+})
+
+test("computeInputBudget returns undefined when context is unknown", () => {
+    assert.equal(computeInputBudget({ context: 0, input: 100, output: 50 }), undefined)
+})
+
+test("computeInputBudget never returns negative when output exceeds context", () => {
+    assert.equal(computeInputBudget({ context: 100, output: 200 }), 0)
+})
+
+test("computeInputBudget prefers explicit input over the context-minus-output fallback", () => {
+    // If both `input` and `output` are present, `input` wins regardless of what
+    // `context - output` would compute to. Defensive against models where the
+    // numbers don't satisfy `input + output = context`.
+    assert.equal(computeInputBudget({ context: 1000, input: 500, output: 200 }), 500)
+})