From f2fa575a5221f361371980cf06ebdd5d728ce63e Mon Sep 17 00:00:00 2001
From: AgentSeal <hello@agentseal.org>
Date: Wed, 8 Jul 2026 22:55:44 +0200
Subject: [PATCH] Fix model pricing variant resolution

---
 src/models.ts        | 26 +++++++++++++++++++++
 tests/models.test.ts | 55 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/src/models.ts b/src/models.ts
index a6c789b..7586fcd 100644
--- a/src/models.ts
+++ b/src/models.ts
@@ -289,6 +289,8 @@ const BUILTIN_ALIASES: Record<string, string> = {
   'kimi-auto':                     'kimi-k2-thinking',
   'kimi-code':                     'kimi-k2-thinking',
   'kimi-for-coding':               'kimi-k2-thinking',
+  'mimo-v2-flash':                 'xiaomi/mimo-v2-flash',
+  'kat-coder-pro-v1':              'kwaipilot/kat-coder-pro',
   // Cursor emits dot-version tier-last names plus tier/reasoning suffixes
   // that LiteLLM does not index (`-high`, `-low`, `-medium`, `-thinking`,
   // `-high-thinking`, `-fast-mode`). Missing aliases here surface as $0 in
@@ -575,6 +577,8 @@ export function getProxyPathsConfigHash(): string {
 function resolveAlias(model: string): string {
   if (Object.hasOwn(userAliases, model)) return userAliases[model]!
   if (Object.hasOwn(BUILTIN_ALIASES, model)) return BUILTIN_ALIASES[model]!
+  const lowercase = model.toLowerCase()
+  if (lowercase !== model && Object.hasOwn(BUILTIN_ALIASES, lowercase)) return BUILTIN_ALIASES[lowercase]!
   return model
 }
 function getCanonicalName(model: string): string {
@@ -584,6 +588,16 @@ function getCanonicalName(model: string): string {
     .replace(/^[^/]+\//, '') // strip provider prefix: anthropic/foo -> foo
 }
 
+function stripKnownPricingVariantSuffix(model: string): string | null {
+  const withoutColonSuffix = model.replace(/:(thinking|cloud)$/i, '')
+  if (withoutColonSuffix !== model) return withoutColonSuffix
+
+  const withoutTeeSuffix = model.replace(/-TEE$/i, '')
+  if (withoutTeeSuffix !== model) return withoutTeeSuffix
+
+  return null
+}
+
 export function getModelCosts(model: string): ModelCosts | null {
   // Try with provider prefix preserved (azure/gpt-5.4, openrouter/anthropic/claude-opus-4.6)
   const withPrefix = model.replace(/@.*$/, '').replace(/-\d{8}$/, '')
@@ -631,6 +645,18 @@ export function getModelCosts(model: string): ModelCosts | null {
   const byPrefix = lowerIndex.get(withPrefix.toLowerCase())
   if (byPrefix) return byPrefix
 
+  const withPrefixVariant = stripKnownPricingVariantSuffix(withPrefix)
+  if (withPrefixVariant && withPrefixVariant !== withPrefix) {
+    const variantCosts = getModelCosts(withPrefixVariant)
+    if (variantCosts) return variantCosts
+  }
+
+  const canonicalVariant = stripKnownPricingVariantSuffix(canonical)
+  if (canonicalVariant && canonicalVariant !== canonical && canonicalVariant !== withPrefixVariant) {
+    const variantCosts = getModelCosts(canonicalVariant)
+    if (variantCosts) return variantCosts
+  }
+
   return null
 }
 
diff --git a/tests/models.test.ts b/tests/models.test.ts
index 924b914..4ce651f 100644
--- a/tests/models.test.ts
+++ b/tests/models.test.ts
@@ -620,3 +620,58 @@ describe('DeepSeek v4 models resolve to pricing', () => {
     }
   })
 })
+
+describe('provider pricing suffix variants', () => {
+  const cases: Array<[string, string]> = [
+    ['GLM-4.7-TEE', 'glm-4.7'],
+    ['glm-4.7:thinking', 'glm-4.7'],
+    ['Kimi-K2.5-TEE', 'kimi-k2.5'],
+    ['deepseek-v4-pro:cloud', 'deepseek-v4-pro'],
+    ['glm-5:thinking', 'glm-5'],
+    ['kimi-k2.6:thinking', 'kimi-k2.6'],
+    ['deepseek-v4-flash:thinking', 'deepseek-v4-flash'],
+    ['minimax-m3:cloud', 'minimax-m3'],
+  ]
+
+  for (const [input, expectedBase] of cases) {
+    it(`${input} resolves through ${expectedBase}`, () => {
+      const costs = getModelCosts(input)
+      const expected = getModelCosts(expectedBase)
+      expect(costs).not.toBeNull()
+      expect(expected).not.toBeNull()
+      expect(costs!.inputCostPerToken).toBe(expected!.inputCostPerToken)
+      expect(costs!.outputCostPerToken).toBe(expected!.outputCostPerToken)
+    })
+  }
+
+  it('does not strip arbitrary local runtime tags', () => {
+    expect(getModelCosts('qwen3.6:35b-a3b-bf16')).toBeNull()
+  })
+
+  it('does not strip free-tier markers into paid pricing', () => {
+    expect(getModelCosts('mimo-v2-flash:free')).toBeNull()
+  })
+})
+
+describe('observed provider model aliases', () => {
+  const cases: Array<[string, string]> = [
+    ['MiMo-V2-Flash', 'xiaomi/mimo-v2-flash'],
+    ['KAT-Coder-Pro-V1', 'kwaipilot/kat-coder-pro'],
+  ]
+
+  for (const [input, expectedModel] of cases) {
+    it(`${input} resolves through ${expectedModel}`, () => {
+      const costs = getModelCosts(input)
+      const expected = getModelCosts(expectedModel)
+      expect(costs).not.toBeNull()
+      expect(expected).not.toBeNull()
+      expect(costs).toEqual(expected)
+      expect(calculateCost(input, 1_000_000, 1_000_000, 0, 0, 0)).toBeGreaterThan(0)
+    })
+  }
+
+  it('does not map dated Qwen3 Max to a reseller price without provider context', () => {
+    expect(getModelCosts('qwen3-max-2026-01-23')).toBeNull()
+    expect(calculateCost('qwen3-max-2026-01-23', 1_000_000, 1_000_000, 0, 0, 0)).toBe(0)
+  })
+})