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.
This commit is contained in:
@@ -58,7 +58,7 @@ const EMPTY_COMPILED_EQUIVALENCE: CompiledEquivalenceConfig = {
|
||||
};
|
||||
const kModelResolutionCache = Symbol("model-equivalence.resolutionCache");
|
||||
interface CompiledEquivalenceConfigWithCache extends CompiledEquivalenceConfig {
|
||||
[kModelResolutionCache]?: WeakMap<Model<Api>, ResolvedCanonicalModel>;
|
||||
[kModelResolutionCache]?: Map<string, ResolvedCanonicalModel>;
|
||||
}
|
||||
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<Api>): string {
|
||||
return `${model.provider}/${model.id}`;
|
||||
}
|
||||
|
||||
function getModelResolutionCacheKey(model: Model<Api>): string {
|
||||
return `${model.provider}\0${model.id}`;
|
||||
}
|
||||
|
||||
function buildOverrideMap(overrides: Record<string, string> | undefined): Map<string, string> {
|
||||
const result = new Map<string, string>();
|
||||
if (!overrides) {
|
||||
@@ -728,10 +736,10 @@ function getPreferredFallbackCanonicalCandidate(modelId: string, candidates: rea
|
||||
|
||||
function resolveCanonicalIdForModel(
|
||||
model: Model<Api>,
|
||||
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<Model<Api>, ResolvedCanonicalModel>();
|
||||
modelCache = new Map<string, ResolvedCanonicalModel>();
|
||||
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,
|
||||
|
||||
@@ -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<string>();
|
||||
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<string>();
|
||||
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<string>();
|
||||
for (const candidate of [withoutBoth, withoutLeading, withoutTrailing]) {
|
||||
if (candidate && candidate !== normalized) {
|
||||
candidates.add(candidate);
|
||||
|
||||
@@ -37,6 +37,9 @@ export interface ContextBreakdown {
|
||||
freeTokens: number;
|
||||
}
|
||||
|
||||
const EMPTY_STRING_PARTS: readonly string[] = [];
|
||||
const EMPTY_TOOLS: ReadonlyArray<Pick<Tool, "name" | "description" | "parameters">> = [];
|
||||
|
||||
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;
|
||||
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -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<skills>\n- code: Write code\n- review: Review code\n</skills>",
|
||||
"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);
|
||||
|
||||
Reference in New Issue
Block a user