diff --git a/src/__tests__/models.test.ts b/src/__tests__/models.test.ts new file mode 100644 index 0000000..b8f83e8 --- /dev/null +++ b/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() + } + }) +}) diff --git a/src/core/request/request-handler.ts b/src/core/request/request-handler.ts index e7745b7..5bf3f67 100644 --- a/src/core/request/request-handler.ts +++ b/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 diff --git a/src/plugin/models.ts b/src/plugin/models.ts index c9aac7a..67d7524 100644 --- a/src/plugin/models.ts +++ b/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 +} + +const catalogCache = new Map() +let activeContextWindows = new Map() 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): void { + activeContextWindows = new Map(contextWindows) +} + +function parseModelCatalog(data: ModelCatalogResponse): Map { + const contextWindows = new Map() + + 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 { + 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() }