From fde55bf927d059dfd638c8bcefca6ecbccb3f5f8 Mon Sep 17 00:00:00 2001 From: can1357 Date: Sat, 6 Jun 2026 22:21:58 +0200 Subject: [PATCH] fix(model): added bracket-affix stripping and string-keyed resolution cache - Replaced WeakMap model cache with provider/id string keys for stable reuse. - Returned official model ids directly when matched, before heuristics. - Collapsed non-message token path to system prompt and tool schema totals. --- .../src/config/model-equivalence.ts | 31 ++++++--- .../src/config/model-id-affixes.ts | 61 ++++++++++------- .../src/modes/utils/context-usage.ts | 16 +++-- .../test/model-id-affixes.test.ts | 66 +++++++++++++++++++ .../test/status-line-context-cache.test.ts | 34 ++++++++++ 5 files changed, 171 insertions(+), 37 deletions(-) create mode 100644 packages/coding-agent/test/model-id-affixes.test.ts diff --git a/packages/coding-agent/src/config/model-equivalence.ts b/packages/coding-agent/src/config/model-equivalence.ts index 28ae3744a..e30755b2d 100644 --- a/packages/coding-agent/src/config/model-equivalence.ts +++ b/packages/coding-agent/src/config/model-equivalence.ts @@ -58,7 +58,7 @@ const EMPTY_COMPILED_EQUIVALENCE: CompiledEquivalenceConfig = { }; const kModelResolutionCache = Symbol("model-equivalence.resolutionCache"); interface CompiledEquivalenceConfigWithCache extends CompiledEquivalenceConfig { - [kModelResolutionCache]?: WeakMap, ResolvedCanonicalModel>; + [kModelResolutionCache]?: Map; } const FAMILY_EXTRACTION_PATTERNS = [ /(?:^|[/:._-])((?:claude|gemini|gpt|grok|glm|qwen|minimax|kimi|deepseek|llama|gemma|nova|mistral|ministral|pixtral|codestral|devstral|magistral|ernie|doubao|seed|aion|olmo|molmo|nemotron|palmyra|command|codex|coder|o[1345])[-a-z0-9.]+)(?::|$)/i, @@ -128,10 +128,18 @@ function normalizeCanonicalIdKey(canonicalId: string): string { return canonicalId.trim().toLowerCase(); } +function getCanonicalSuffixAliasKey(candidate: string): string { + return PENALTY_HAS_UPPERCASE.test(candidate) ? normalizeCanonicalIdKey(candidate) : candidate; +} + export function formatCanonicalVariantSelector(model: Model): string { return `${model.provider}/${model.id}`; } +function getModelResolutionCacheKey(model: Model): string { + return `${model.provider}\0${model.id}`; +} + function buildOverrideMap(overrides: Record | undefined): Map { const result = new Map(); if (!overrides) { @@ -728,10 +736,10 @@ function getPreferredFallbackCanonicalCandidate(modelId: string, candidates: rea function resolveCanonicalIdForModel( model: Model, + selector: string, equivalence: CompiledEquivalenceConfig, referenceData: CanonicalReferenceData, ): ResolvedCanonicalModel { - const selector = formatCanonicalVariantSelector(model); const normalizedSelector = normalizeSelectorKey(selector); if (equivalence.overrides.has(normalizedSelector)) { @@ -752,10 +760,14 @@ function resolveCanonicalIdForModel( return { id: claudeFamilyAlias, source: claudeFamilyAlias === model.id ? "bundled" : "heuristic" }; } + if (referenceData.officialIds.has(model.id) && !model.id.includes("/") && !model.id.includes(":")) { + return { id: model.id, source: "bundled" }; + } + const heuristicCandidates = getHeuristicCanonicalCandidates(model.id, referenceData.officialIds); const officialMatches = new Set(heuristicCandidates.filter(candidate => referenceData.officialIds.has(candidate))); for (const candidate of heuristicCandidates) { - const aliased = referenceData.suffixAliases.get(normalizeCanonicalIdKey(candidate)); + const aliased = referenceData.suffixAliases.get(getCanonicalSuffixAliasKey(candidate)); if (aliased) { officialMatches.add(aliased); } @@ -814,17 +826,18 @@ export function buildCanonicalModelIndex( const compiledWithCache = compiledEquivalence as CompiledEquivalenceConfigWithCache; let modelCache = compiledWithCache[kModelResolutionCache]; if (!modelCache) { - modelCache = new WeakMap, ResolvedCanonicalModel>(); + modelCache = new Map(); compiledWithCache[kModelResolutionCache] = modelCache; } for (const model of models) { - let canonical = modelCache.get(model); - if (!canonical) { - canonical = resolveCanonicalIdForModel(model, compiledEquivalence, referenceData); - modelCache.set(model, canonical); - } const selector = formatCanonicalVariantSelector(model); + const cacheKey = getModelResolutionCacheKey(model); + let canonical = modelCache.get(cacheKey); + if (!canonical) { + canonical = resolveCanonicalIdForModel(model, selector, compiledEquivalence, referenceData); + modelCache.set(cacheKey, canonical); + } const variant: CanonicalModelVariant = { canonicalId: canonical.id, selector, diff --git a/packages/coding-agent/src/config/model-id-affixes.ts b/packages/coding-agent/src/config/model-id-affixes.ts index b4aff136f..7cec31892 100644 --- a/packages/coding-agent/src/config/model-id-affixes.ts +++ b/packages/coding-agent/src/config/model-id-affixes.ts @@ -4,34 +4,49 @@ const MODEL_ID_SEGMENT_PATTERN = /[a-z0-9.:-]+/g; const MODEL_FAMILY_PREFIX_PATTERN = /^(claude|gemini|gpt|grok|glm|qwen|deepseek|kimi|mimo|doubao|ernie|gpt-oss|gemma|minimax|step|command|jamba|llama|o[1345])/i; -function hasDigit(value: string): boolean { - return /\d/.test(value); +function normalizeModelIdWhitespace(value: string): string { + return value.trim().replace(/\s+/g, " "); } +/** Ordering for model-like segments: longest first, ties broken lexicographically. */ function compareSegmentPreference(left: string, right: string): number { - if (left.length !== right.length) { - return right.length - left.length; - } - return left.localeCompare(right); + return left.length !== right.length ? right.length - left.length : left.localeCompare(right); } export function getModelLikeIdSegments(modelId: string): string[] { - const normalized = normalizeModelIdWhitespace(modelId).toLowerCase(); - if (!normalized) return []; - const segments = (normalized.match(MODEL_ID_SEGMENT_PATTERN) ?? []).filter( - segment => MODEL_FAMILY_PREFIX_PATTERN.test(segment) && hasDigit(segment), - ); - const unique = [...new Set(segments)]; - unique.sort(compareSegmentPreference); - return unique; + const matches = normalizeModelIdWhitespace(modelId).toLowerCase().match(MODEL_ID_SEGMENT_PATTERN); + if (!matches) return []; + const segments = new Set(); + for (const segment of matches) { + if (MODEL_FAMILY_PREFIX_PATTERN.test(segment) && /\d/.test(segment)) segments.add(segment); + } + return [...segments].sort(compareSegmentPreference); } export function getLongestModelLikeIdSegment(modelId: string): string | undefined { - return getModelLikeIdSegments(modelId)[0]; + const matches = normalizeModelIdWhitespace(modelId).toLowerCase().match(MODEL_ID_SEGMENT_PATTERN); + if (!matches) return undefined; + let best: string | undefined; + for (const segment of matches) { + if ( + MODEL_FAMILY_PREFIX_PATTERN.test(segment) && + /\d/.test(segment) && + (best === undefined || compareSegmentPreference(segment, best) < 0) + ) { + best = segment; + } + } + return best; } -function normalizeModelIdWhitespace(value: string): string { - return value.trim().replace(/\s+/g, " "); +function hasBracketAffixMarker(value: string): boolean { + for (let index = 0; index < value.length; index++) { + const code = value.charCodeAt(index); + if (code === 91 || code === 93 || code === 0x3010 || code === 0x3011) { + return true; + } + } + return false; } /** @@ -39,18 +54,20 @@ function normalizeModelIdWhitespace(value: string): string { * upstream model id, e.g. * "[Kiro] claude-opus-4-8" -> "claude-opus-4-8" * "[gcli转] gemini-3.1-pro-preview [假流]" -> "gemini-3.1-pro-preview" + * + * Candidates are returned most-stripped first: both ends, then leading-only, then trailing-only. */ export function getBracketStrippedModelIdCandidates(modelId: string): string[] { + if (!hasBracketAffixMarker(modelId)) return []; const normalized = normalizeModelIdWhitespace(modelId); if (!normalized) return []; - const candidates = new Set(); - const withoutLeading = normalizeModelIdWhitespace(normalized.replace(LEADING_BRACKETED_AFFIX_PATTERN, "")); + const strippedLeading = normalized.replace(LEADING_BRACKETED_AFFIX_PATTERN, ""); + const withoutLeading = normalizeModelIdWhitespace(strippedLeading); const withoutTrailing = normalizeModelIdWhitespace(normalized.replace(TRAILING_BRACKETED_AFFIX_PATTERN, "")); - const withoutBoth = normalizeModelIdWhitespace( - normalized.replace(LEADING_BRACKETED_AFFIX_PATTERN, "").replace(TRAILING_BRACKETED_AFFIX_PATTERN, ""), - ); + const withoutBoth = normalizeModelIdWhitespace(strippedLeading.replace(TRAILING_BRACKETED_AFFIX_PATTERN, "")); + const candidates = new Set(); for (const candidate of [withoutBoth, withoutLeading, withoutTrailing]) { if (candidate && candidate !== normalized) { candidates.add(candidate); diff --git a/packages/coding-agent/src/modes/utils/context-usage.ts b/packages/coding-agent/src/modes/utils/context-usage.ts index fd93070a8..d223d9a5e 100644 --- a/packages/coding-agent/src/modes/utils/context-usage.ts +++ b/packages/coding-agent/src/modes/utils/context-usage.ts @@ -37,6 +37,9 @@ export interface ContextBreakdown { freeTokens: number; } +const EMPTY_STRING_PARTS: readonly string[] = []; +const EMPTY_TOOLS: ReadonlyArray> = []; + export function estimateSkillsTokens(skills: readonly Skill[]): number { const fragments: string[] = []; for (const skill of skills) { @@ -75,15 +78,16 @@ export function estimateToolSchemaTokens( * messages walked incrementally as new entries append. */ export function computeNonMessageTokens(session: AgentSession): number { - const parts = computeNonMessageBreakdown(session); - return parts.systemPromptTokens + parts.systemContextTokens + parts.toolsTokens + parts.skillsTokens; + const systemPromptParts = session.systemPrompt ?? EMPTY_STRING_PARTS; + const tools = session.agent?.state?.tools ?? EMPTY_TOOLS; + return countTokens(systemPromptParts) + estimateToolSchemaTokens(tools); } /** - * Shared helper for the four non-message token totals. Single source of truth - * for both `computeNonMessageTokens` (status-line incremental cache) and - * `computeContextBreakdown` (/context panel). The split avoids drift between - * the two surfaces — they MUST report the same numbers. + * Shared helper for the four non-message token totals used by + * `computeContextBreakdown` (/context panel). Keep this category split stable: + * the status-line fast path intentionally uses the equivalent collapsed total + * in `computeNonMessageTokens`. */ function computeNonMessageBreakdown(session: AgentSession): { skillsTokens: number; diff --git a/packages/coding-agent/test/model-id-affixes.test.ts b/packages/coding-agent/test/model-id-affixes.test.ts new file mode 100644 index 000000000..b94a9bc48 --- /dev/null +++ b/packages/coding-agent/test/model-id-affixes.test.ts @@ -0,0 +1,66 @@ +import { describe, expect, test } from "bun:test"; +import { + getBracketStrippedModelIdCandidates, + getLongestModelLikeIdSegment, + getModelLikeIdSegments, + stripBracketedModelIdAffixes, +} from "../src/config/model-id-affixes"; + +describe("getModelLikeIdSegments", () => { + test("keeps only family-prefixed segments that carry a digit, deduped", () => { + expect(getModelLikeIdSegments("openrouter/anthropic/claude-3.5-sonnet")).toEqual(["claude-3.5-sonnet"]); + // `random-text` lacks a family prefix; `claude` (no digit) is dropped. + expect(getModelLikeIdSegments("random-text claude gemini-2")).toEqual(["gemini-2"]); + }); + + test("orders longest first with lexicographic tie-break", () => { + expect(getModelLikeIdSegments("claude-3 claude-3-5-haiku claude-2")).toEqual([ + "claude-3-5-haiku", + "claude-2", + "claude-3", + ]); + }); + + test("normalizes whitespace and case before matching", () => { + expect(getModelLikeIdSegments(" GLM-4.5-Air GEMINI-2 ")).toEqual(["glm-4.5-air", "gemini-2"]); + }); + + test("returns empty for ids with no model-like segment", () => { + expect(getModelLikeIdSegments("")).toEqual([]); + expect(getModelLikeIdSegments("just some words")).toEqual([]); + }); +}); + +describe("getLongestModelLikeIdSegment", () => { + test("matches getModelLikeIdSegments[0]", () => { + const id = "[Kiro] claude-3 claude-3-5-sonnet"; + expect(getLongestModelLikeIdSegment(id)).toBe(getModelLikeIdSegments(id)[0]); + expect(getLongestModelLikeIdSegment(id)).toBe("claude-3-5-sonnet"); + }); + + test("is undefined when nothing matches", () => { + expect(getLongestModelLikeIdSegment("vendor/unknown-tag")).toBeUndefined(); + }); +}); + +describe("getBracketStrippedModelIdCandidates", () => { + test("no brackets yields no candidates", () => { + expect(getBracketStrippedModelIdCandidates("claude-opus-4-8")).toEqual([]); + }); + + test("strips leading reseller tag", () => { + expect(getBracketStrippedModelIdCandidates("[Kiro] claude-opus-4-8")).toEqual(["claude-opus-4-8"]); + }); + + test("strips both ends first, then each side, in preference order", () => { + expect(getBracketStrippedModelIdCandidates("[gcli转] gemini-3.1-pro-preview [假流]")).toEqual([ + "gemini-3.1-pro-preview", + "gemini-3.1-pro-preview [假流]", + "[gcli转] gemini-3.1-pro-preview", + ]); + }); + + test("supports full-width brackets", () => { + expect(stripBracketedModelIdAffixes("【供应商】 deepseek-v3 【限时】")).toBe("deepseek-v3"); + }); +}); diff --git a/packages/coding-agent/test/status-line-context-cache.test.ts b/packages/coding-agent/test/status-line-context-cache.test.ts index 7abacdfdd..95ff25dc5 100644 --- a/packages/coding-agent/test/status-line-context-cache.test.ts +++ b/packages/coding-agent/test/status-line-context-cache.test.ts @@ -15,8 +15,10 @@ * (messages.length shrinks) resets the cache. */ import { afterAll, beforeAll, describe, expect, it } from "bun:test"; +import { countTokens } from "@oh-my-pi/pi-natives"; import { resetSettingsForTest, Settings } from "../src/config/settings"; import { StatusLineComponent } from "../src/modes/components/status-line"; +import { computeNonMessageTokens, estimateToolSchemaTokens } from "../src/modes/utils/context-usage"; import { initTheme } from "../src/modes/theme/theme"; import type { AgentSession } from "../src/session/agent-session"; @@ -122,6 +124,38 @@ describe("StatusLineComponent incremental context breakdown cache", () => { expect(v3.usedTokens).toBeGreaterThan(v2.usedTokens); }); + it("non-message token shortcut matches previous category sum semantics", () => { + const session = makeSession({ + messages: [], + systemPrompt: [ + "You are an assistant.\n\n\n- code: Write code\n- review: Review code\n", + "Loaded context file", + "Runtime note", + ], + tools: [ + { + name: "bash", + description: "Run shell commands", + parameters: { type: "object", properties: { command: { type: "string" } } }, + }, + ], + skills: [ + { name: "code", description: "Write code" }, + { name: "review", description: "Review code" }, + ], + }); + + const skillsTokens = countTokens(["code", "Write code", "review", "Review code"]); + const previousCategorySum = + Math.max(0, countTokens(session.systemPrompt?.[0] ?? "") - skillsTokens) + + countTokens((session.systemPrompt ?? []).slice(1)) + + estimateToolSchemaTokens(session.agent?.state?.tools ?? []) + + skillsTokens; + + expect(new StatusLineComponent(session).getCachedContextBreakdown().usedTokens).toBe(previousCategorySum); + expect(computeNonMessageTokens(session)).toBe(previousCategorySum); + }); + it("zero messages: produces only non-message tokens, no crash", () => { const session = makeSession({ messages: [] }); const comp = new StatusLineComponent(session);