feat(models): discover context windows from Kiro catalog

src/__tests__/models.test.ts

@@ -0,0 +1,40 @@
+import { describe, expect, test } from 'bun:test'
+import {
+  clearContextWindowCatalog,
+  getContextWindowSize,
+  refreshContextWindowSizes
+} from '../plugin/models.js'
+import type { KiroAuthDetails } from '../plugin/types.js'
+
+const auth: KiroAuthDetails = {
+  access: 'test-token',
+  refresh: 'test-refresh',
+  expires: Date.now() + 60_000,
+  authMethod: 'idc',
+  region: 'us-east-1'
+}
+
+describe('getContextWindowSize', () => {
+  test('uses maxInputTokens from the live model catalog for aliases', async () => {
+    const originalFetch = globalThis.fetch
+    globalThis.fetch = (async () =>
+      new Response(
+        JSON.stringify({
+          models: [
+            { modelId: 'claude-sonnet-4.6', tokenLimits: { maxInputTokens: 1_000_000 } },
+            { modelId: 'deepseek-3.2', tokenLimits: { maxInputTokens: 164_000 } }
+          ]
+        }),
+        { status: 200 }
+      )) as typeof fetch
+
+    try {
+      await refreshContextWindowSizes(auth)
+      expect(getContextWindowSize('claude-sonnet-4-6')).toBe(1_000_000)
+      expect(getContextWindowSize('deepseek-3.2')).toBe(164_000)
+    } finally {
+      globalThis.fetch = originalFetch
+      clearContextWindowCatalog()
+    }
+  })
+})

src/core/request/request-handler.ts

@@ -4,6 +4,7 @@ import type { AccountManager } from '../../plugin/accounts'
 import type { KiroConfig } from '../../plugin/config'
 import { isPermanentError } from '../../plugin/health'
 import * as logger from '../../plugin/logger'
+import { refreshContextWindowSizes } from '../../plugin/models'
 import { transformToSdkRequest } from '../../plugin/request'
 import { createSdkClient } from '../../plugin/sdk-client'
 import { syncFromKiroCli } from '../../plugin/sync/kiro-cli'
@@ -132,6 +133,7 @@ export class RequestHandler {
         continue
       }
 
+      await refreshContextWindowSizes(auth)
       const sdkPrep = this.prepareSdkRequest(init?.body, model, auth, think, budget, showToast)
 
       const apiTimestamp = this.config.enable_log_api_request ? logger.getTimestamp() : null

src/plugin/models.ts

@@ -1,4 +1,28 @@
-import { MODEL_MAPPING, SUPPORTED_MODELS, isLongContextModel } from '../constants'
+import {
+  extractRegionFromArn,
+  isLongContextModel,
+  MODEL_MAPPING,
+  SUPPORTED_MODELS
+} from '../constants'
+import type { KiroAuthDetails } from './types'
+
+const MODEL_CATALOG_TTL_MS = 5 * 60 * 1000
+const DEFAULT_CONTEXT_WINDOW = 200_000
+
+type ModelCatalogResponse = {
+  models?: Array<{
+    modelId?: string
+    tokenLimits?: { maxInputTokens?: number }
+  }>
+}
+
+type CatalogCacheEntry = {
+  expiresAt: number
+  contextWindows: Map<string, number>
+}
+
+const catalogCache = new Map<string, CatalogCacheEntry>()
+let activeContextWindows = new Map<string, number>()
 
 export function resolveKiroModel(model: string): string {
   const resolved = MODEL_MAPPING[model]
@@ -9,5 +33,85 @@ export function resolveKiroModel(model: string): string {
 }
 
 export function getContextWindowSize(model: string): number {
-  return isLongContextModel(model) ? 1000000 : 200000
+  return (
+    activeContextWindows.get(model) ??
+    (isLongContextModel(model) ? 1_000_000 : DEFAULT_CONTEXT_WINDOW)
+  )
+}
+
+function getCatalogRegion(auth: KiroAuthDetails): string {
+  return extractRegionFromArn(auth.profileArn) ?? auth.region ?? 'us-east-1'
+}
+
+function applyCatalog(contextWindows: Map<string, number>): void {
+  activeContextWindows = new Map(contextWindows)
+}
+
+function parseModelCatalog(data: ModelCatalogResponse): Map<string, number> {
+  const contextWindows = new Map<string, number>()
+
+  for (const model of data.models ?? []) {
+    const modelId = model.modelId
+    const maxInputTokens = model.tokenLimits?.maxInputTokens
+    if (
+      !modelId ||
+      typeof maxInputTokens !== 'number' ||
+      !Number.isInteger(maxInputTokens) ||
+      maxInputTokens <= 0
+    )
+      continue
+
+    contextWindows.set(modelId, maxInputTokens)
+    for (const [alias, resolved] of Object.entries(MODEL_MAPPING)) {
+      if (resolved === modelId) contextWindows.set(alias, maxInputTokens)
+    }
+  }
+
+  return contextWindows
+}
+
+export async function refreshContextWindowSizes(auth: KiroAuthDetails): Promise<void> {
+  const cacheKey = `${auth.access}:${getCatalogRegion(auth)}`
+  const cached = catalogCache.get(cacheKey)
+  if (cached && cached.expiresAt > Date.now()) {
+    applyCatalog(cached.contextWindows)
+    return
+  }
+  activeContextWindows = new Map()
+
+  const regions = [getCatalogRegion(auth)]
+  if (!regions.includes('us-east-1')) regions.push('us-east-1')
+
+  for (const region of regions) {
+    const endpoint = new URL(`https://q.${region}.amazonaws.com/ListAvailableModels`)
+    endpoint.searchParams.set('origin', 'AI_EDITOR')
+
+    try {
+      const response = await fetch(endpoint, {
+        headers: {
+          Accept: 'application/json',
+          Authorization: `Bearer ${auth.access}`
+        },
+        signal: AbortSignal.timeout(5_000)
+      })
+      if (!response.ok) continue
+
+      const contextWindows = parseModelCatalog((await response.json()) as ModelCatalogResponse)
+      if (contextWindows.size === 0) continue
+
+      catalogCache.set(cacheKey, {
+        expiresAt: Date.now() + MODEL_CATALOG_TTL_MS,
+        contextWindows
+      })
+      applyCatalog(contextWindows)
+      return
+    } catch {
+      continue
+    }
+  }
+}
+
+export function clearContextWindowCatalog(): void {
+  catalogCache.clear()
+  activeContextWindows = new Map()
 }