From 89d64de8860b091ce712e61a21dd5410262b7beb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E5=9F=BA=E9=AD=81?= <1412414664@qq.com> Date: Mon, 6 Jul 2026 20:34:28 +0800 Subject: [PATCH 1/3] fix(model): keep callback data short --- .../model-selection-callback-handler.ts | 160 ++++++++++++++---- src/bot/menus/model-selection-menu.ts | 26 ++- tests/bot/handlers/model.test.ts | 160 +++++++++++++++++- 3 files changed, 303 insertions(+), 43 deletions(-) diff --git a/src/bot/callbacks/model-selection-callback-handler.ts b/src/bot/callbacks/model-selection-callback-handler.ts index a926e59c..3903e6f4 100644 --- a/src/bot/callbacks/model-selection-callback-handler.ts +++ b/src/bot/callbacks/model-selection-callback-handler.ts @@ -1,6 +1,7 @@ import { Context, InlineKeyboard } from "grammy"; import { getStoredAgent, resolveProjectAgent } from "../../app/services/agent-selection-service.js"; import { + getModelSelectionLists, searchModels, selectModel, } from "../../app/services/model-selection-service.js"; @@ -15,15 +16,46 @@ import { keyboardManager } from "../keyboards/keyboard-manager.js"; import { pinnedMessageManager } from "../pinned/pinned-message-manager.js"; import { clearActiveInlineMenu, ensureActiveInlineMenu } from "../menus/inline-menu.js"; import { + MODEL_LIST_CALLBACK_PREFIX, MODEL_SEARCH_AGAIN_CALLBACK, MODEL_SEARCH_CALLBACK, MODEL_SEARCH_CANCEL_CALLBACK, } from "../menus/model-selection-menu.js"; +const MODEL_SEARCH_RESULT_CALLBACK_PREFIX = "model:result:"; + interface ModelSearchMetadata { flow: string; stage: string; messageId?: number; + models: ModelInfo[]; +} + +function parseModelItems(value: unknown): ModelInfo[] { + if (!Array.isArray(value)) { + return []; + } + + return value.flatMap((item) => { + if ( + typeof item !== "object" || + item === null || + !("providerID" in item) || + !("modelID" in item) + ) { + return []; + } + + const providerID = item.providerID; + const modelID = item.modelID; + if (typeof providerID !== "string" || typeof modelID !== "string") { + return []; + } + + const variant = + "variant" in item && typeof item.variant === "string" ? item.variant : "default"; + return [{ providerID, modelID, variant }]; + }); } function parseModelSearchMetadata(): ModelSearchMetadata | null { @@ -42,17 +74,83 @@ function parseModelSearchMetadata(): ModelSearchMetadata | null { const messageId = typeof state.metadata.messageId === "number" ? state.metadata.messageId : undefined; - return { flow, stage, messageId }; + return { flow, stage, messageId, models: parseModelItems(state.metadata.models) }; +} + +function parseNonNegativeIndex(value: string): number | null { + if (!/^\d+$/.test(value)) { + return null; + } + + const index = Number.parseInt(value, 10); + if (!Number.isInteger(index) || index < 0) { + return null; + } + + return index; +} + +function parseCallbackIndex(data: string, prefix: string): number | null { + if (!data.startsWith(prefix)) { + return null; + } + + return parseNonNegativeIndex(data.slice(prefix.length)); +} + +async function resolveModelListCallback(data: string): Promise { + if (!data.startsWith(MODEL_LIST_CALLBACK_PREFIX)) { + return null; + } + + const parts = data.slice(MODEL_LIST_CALLBACK_PREFIX.length).split(":"); + if (parts.length !== 2) { + return null; + } + + const [kind, indexText] = parts; + const index = parseNonNegativeIndex(indexText); + if ((kind !== "favorites" && kind !== "recent") || index === null) { + return null; + } + + const lists = await getModelSelectionLists(); + const model = kind === "favorites" ? lists.favorites[index] : lists.recent[index]; + if (!model) { + return null; + } + + return { + providerID: model.providerID, + modelID: model.modelID, + variant: "default", + }; +} + +function parseLegacyModelCallback(data: string): ModelInfo | null { + const parts = data.split(":"); + if (parts.length < 3) { + return null; + } + + const providerID = parts[1]; + const modelID = parts.slice(2).join(":"); + if (!providerID || !modelID) { + return null; + } + + return { + providerID, + modelID, + variant: "default", + }; } /** * Shared logic for applying a model selection and updating UI. * Used by both the regular inline menu flow and the search results flow. */ -async function applyModelSelectionAndNotify( - ctx: Context, - modelInfo: ModelInfo, -): Promise { +async function applyModelSelectionAndNotify(ctx: Context, modelInfo: ModelInfo): Promise { if (ctx.chat) { keyboardManager.initialize(ctx.api, ctx.chat.id); } @@ -119,24 +217,17 @@ export async function handleModelSelect(ctx: Context): Promise { logger.debug(`[ModelHandler] Received callback: ${callbackQuery.data}`); try { - // Parse callback data: "model:providerID:modelID" - const parts = callbackQuery.data.split(":"); - if (parts.length < 3) { + const modelInfo = + (await resolveModelListCallback(callbackQuery.data)) ?? + parseLegacyModelCallback(callbackQuery.data); + + if (!modelInfo) { logger.error(`[ModelHandler] Invalid callback data format: ${callbackQuery.data}`); clearActiveInlineMenu("model_select_invalid_callback"); await ctx.answerCallbackQuery({ text: t("model.change_error_callback") }).catch(() => {}); return true; } - const providerID = parts[1]; - const modelID = parts.slice(2).join(":"); // Handle model IDs that may contain ":" - - const modelInfo: ModelInfo = { - providerID, - modelID, - variant: "default", - }; - clearActiveInlineMenu("model_selected"); await applyModelSelectionAndNotify(ctx, modelInfo); @@ -210,9 +301,9 @@ export async function handleModelSearchTextInput(ctx: Context): Promise const keyboard = new InlineKeyboard(); - for (const model of results) { + for (const [index, model] of results.entries()) { const label = `${model.providerID}/${model.modelID}`; - keyboard.text(label, `model:${model.providerID}:${model.modelID}`).row(); + keyboard.text(label, `${MODEL_SEARCH_RESULT_CALLBACK_PREFIX}${index}`).row(); } keyboard.row(); @@ -233,6 +324,11 @@ export async function handleModelSearchTextInput(ctx: Context): Promise flow: "model-search", stage: "results", messageId: sent.message_id, + models: results.map((model) => ({ + providerID: model.providerID, + modelID: model.modelID, + variant: "default", + })), }, }); @@ -299,21 +395,25 @@ export async function handleModelSearchResults(ctx: Context): Promise { return true; } - // Model selection from search results - if (data.startsWith("model:")) { - const parts = data.split(":"); - if (parts.length < 3) { + const resultIndex = parseCallbackIndex(data, MODEL_SEARCH_RESULT_CALLBACK_PREFIX); + if (resultIndex !== null) { + const modelInfo = meta.models[resultIndex]; + if (!modelInfo) { + await ctx.answerCallbackQuery({ text: t("model.change_error_callback") }).catch(() => {}); return true; } - const providerID = parts[1]; - const modelID = parts.slice(2).join(":"); + interactionManager.clear("model_search_selected"); + await applyModelSelectionAndNotify(ctx, modelInfo); + return true; + } - const modelInfo: ModelInfo = { - providerID, - modelID, - variant: "default", - }; + // Backward compatibility for callbacks from already-rendered search result messages. + if (data.startsWith("model:")) { + const modelInfo = parseLegacyModelCallback(data); + if (!modelInfo) { + return true; + } interactionManager.clear("model_search_selected"); await applyModelSelectionAndNotify(ctx, modelInfo); diff --git a/src/bot/menus/model-selection-menu.ts b/src/bot/menus/model-selection-menu.ts index 1caae3c4..3c0b969c 100644 --- a/src/bot/menus/model-selection-menu.ts +++ b/src/bot/menus/model-selection-menu.ts @@ -3,11 +3,7 @@ import { fetchCurrentModel, getModelSelectionLists, } from "../../app/services/model-selection-service.js"; -import type { - FavoriteModel, - ModelInfo, - ModelSelectionLists, -} from "../../app/types/model.js"; +import type { FavoriteModel, ModelInfo, ModelSelectionLists } from "../../app/types/model.js"; import { logger } from "../../utils/logger.js"; import { t } from "../../i18n/index.js"; import { replyWithInlineMenu } from "./inline-menu.js"; @@ -15,6 +11,13 @@ import { replyWithInlineMenu } from "./inline-menu.js"; export const MODEL_SEARCH_CALLBACK = "model:search"; export const MODEL_SEARCH_AGAIN_CALLBACK = "model:search:again"; export const MODEL_SEARCH_CANCEL_CALLBACK = "model:search:cancel"; +export const MODEL_LIST_CALLBACK_PREFIX = "model:list:"; + +type ModelListKind = "favorites" | "recent"; + +export function buildModelListCallback(kind: ModelListKind, index: number): string { + return `${MODEL_LIST_CALLBACK_PREFIX}${kind}:${index}`; +} function buildModelSelectionMenuText(modelLists: ModelSelectionLists): string { const lines = [t("model.menu.select"), t("model.menu.favorites_title")]; @@ -52,7 +55,12 @@ export async function buildModelSelectionMenu( return keyboard; } - const addButton = (model: FavoriteModel, prefix: string): void => { + const addButton = ( + model: FavoriteModel, + prefix: string, + kind: ModelListKind, + index: number, + ): void => { const isActive = currentModel && model.providerID === currentModel.providerID && @@ -61,11 +69,11 @@ export async function buildModelSelectionMenu( const label = `${prefix} ${model.providerID}/${model.modelID}`; const labelWithCheck = isActive ? `✅ ${label}` : label; - keyboard.text(labelWithCheck, `model:${model.providerID}:${model.modelID}`).row(); + keyboard.text(labelWithCheck, buildModelListCallback(kind, index)).row(); }; - favorites.forEach((model) => addButton(model, "⭐")); - recent.forEach((model) => addButton(model, "🕘")); + favorites.forEach((model, index) => addButton(model, "⭐", "favorites", index)); + recent.forEach((model, index) => addButton(model, "🕘", "recent", index)); return keyboard; } diff --git a/tests/bot/handlers/model.test.ts b/tests/bot/handlers/model.test.ts index 7e61f187..563d1e92 100644 --- a/tests/bot/handlers/model.test.ts +++ b/tests/bot/handlers/model.test.ts @@ -9,15 +9,51 @@ const mocked = vi.hoisted(() => ({ interactionManagerTransitionMock: vi.fn(), interactionManagerClearMock: vi.fn(), ensureActiveInlineMenuMock: vi.fn(), + selectModelMock: vi.fn(), + resolveProjectAgentMock: vi.fn(), + keyboardInitializeMock: vi.fn(), + keyboardUpdateModelMock: vi.fn(), + keyboardUpdateAgentMock: vi.fn(), + keyboardUpdateContextMock: vi.fn(), + pinnedRefreshContextLimitMock: vi.fn(), + pinnedGetContextInfoMock: vi.fn(), + pinnedGetContextLimitMock: vi.fn(), + createMainKeyboardMock: vi.fn(), })); vi.mock("../../../src/app/services/model-selection-service.js", () => ({ getModelSelectionLists: mocked.getModelSelectionListsMock, searchModels: mocked.searchModelsMock, - selectModel: vi.fn(), + selectModel: mocked.selectModelMock, fetchCurrentModel: vi.fn(), })); +vi.mock("../../../src/app/services/agent-selection-service.js", () => ({ + getStoredAgent: vi.fn(() => "build"), + resolveProjectAgent: mocked.resolveProjectAgentMock, +})); + +vi.mock("../../../src/bot/keyboards/keyboard-manager.js", () => ({ + keyboardManager: { + initialize: mocked.keyboardInitializeMock, + updateModel: mocked.keyboardUpdateModelMock, + updateAgent: mocked.keyboardUpdateAgentMock, + updateContext: mocked.keyboardUpdateContextMock, + }, +})); + +vi.mock("../../../src/bot/keyboards/main-reply-keyboard.js", () => ({ + createMainKeyboard: mocked.createMainKeyboardMock, +})); + +vi.mock("../../../src/bot/pinned/pinned-message-manager.js", () => ({ + pinnedMessageManager: { + refreshContextLimit: mocked.pinnedRefreshContextLimitMock, + getContextInfo: mocked.pinnedGetContextInfoMock, + getContextLimit: mocked.pinnedGetContextLimitMock, + }, +})); + vi.mock("../../../src/app/managers/interaction-manager.js", () => ({ interactionManager: { getSnapshot: mocked.interactionManagerGetSnapshotMock, @@ -33,11 +69,10 @@ vi.mock("../../../src/bot/menus/inline-menu.js", () => ({ replyWithInlineMenu: vi.fn(), })); -import { - buildModelSelectionMenu, -} from "../../../src/bot/menus/model-selection-menu.js"; +import { buildModelSelectionMenu } from "../../../src/bot/menus/model-selection-menu.js"; import { + handleModelSelect, handleModelSearchCallback, handleModelSearchTextInput, handleModelSearchResults, @@ -64,6 +99,17 @@ describe("bot model selection", () => { mocked.interactionManagerTransitionMock.mockReset(); mocked.interactionManagerClearMock.mockReset(); mocked.ensureActiveInlineMenuMock.mockReset(); + mocked.ensureActiveInlineMenuMock.mockResolvedValue(true); + mocked.selectModelMock.mockReset(); + mocked.resolveProjectAgentMock.mockReset().mockResolvedValue("build"); + mocked.keyboardInitializeMock.mockReset(); + mocked.keyboardUpdateModelMock.mockReset(); + mocked.keyboardUpdateAgentMock.mockReset(); + mocked.keyboardUpdateContextMock.mockReset(); + mocked.pinnedRefreshContextLimitMock.mockReset().mockResolvedValue(undefined); + mocked.pinnedGetContextInfoMock.mockReset().mockReturnValue(null); + mocked.pinnedGetContextLimitMock.mockReset().mockReturnValue(0); + mocked.createMainKeyboardMock.mockReset().mockReturnValue({ keyboard: [["main"]] }); }); describe("buildModelSelectionMenu", () => { @@ -95,6 +141,48 @@ describe("bot model selection", () => { expect(keyboard.inline_keyboard[0][0].text).toBe("🔍 Search"); expect(keyboard.inline_keyboard[0][0].callback_data).toBe("model:search"); }); + + it("uses short callback data for long model IDs", async () => { + const longModelID = "accounts/hubabuba3227-1hvtqlh/deployments/kpwpvuky"; + mocked.getModelSelectionListsMock.mockResolvedValue({ + favorites: [], + recent: [{ providerID: "fireworks", modelID: longModelID }], + }); + + const keyboard = await buildModelSelectionMenu(); + const callbackData = keyboard.inline_keyboard[1][0].callback_data; + + expect(callbackData).toBe("model:list:recent:0"); + expect(Buffer.byteLength(callbackData ?? "", "utf-8")).toBeLessThanOrEqual(64); + expect(callbackData).not.toContain(longModelID); + }); + }); + + describe("handleModelSelect", () => { + it("resolves short list callback data to the original long model ID", async () => { + const longModelID = "accounts/hubabuba3227-1hvtqlh/deployments/kpwpvuky"; + mocked.getModelSelectionListsMock.mockResolvedValue({ + favorites: [], + recent: [{ providerID: "fireworks", modelID: longModelID }], + }); + + const ctx = mockContext({ + callbackQuery: { + data: "model:list:recent:0", + message: { message_id: 999 }, + }, + api: {}, + }); + + const result = await handleModelSelect(ctx); + + expect(result).toBe(true); + expect(mocked.selectModelMock).toHaveBeenCalledWith({ + providerID: "fireworks", + modelID: longModelID, + variant: "default", + }); + }); }); describe("handleModelSearchCallback", () => { @@ -174,6 +262,39 @@ describe("bot model selection", () => { expect(result).toBe(false); }); + + it("uses short callback data for long model IDs in search results", async () => { + const longModelID = "accounts/hubabuba3227-1hvtqlh/deployments/kpwpvuky"; + mocked.interactionManagerGetSnapshotMock.mockReturnValue({ + kind: "custom", + metadata: { flow: "model-search", stage: "input" }, + }); + mocked.searchModelsMock.mockResolvedValue([ + { providerID: "fireworks", modelID: longModelID }, + ]); + + const ctx = mockContext({ + message: { text: "fireworks" }, + }); + + const result = await handleModelSearchTextInput(ctx); + const replyOptions = vi.mocked(ctx.reply).mock.calls[0][1] as { + reply_markup: { inline_keyboard: Array> }; + }; + const callbackData = replyOptions.reply_markup.inline_keyboard[0][0].callback_data; + + expect(result).toBe(true); + expect(callbackData).toBe("model:result:0"); + expect(Buffer.byteLength(callbackData ?? "", "utf-8")).toBeLessThanOrEqual(64); + expect(callbackData).not.toContain(longModelID); + expect(mocked.interactionManagerTransitionMock).toHaveBeenCalledWith( + expect.objectContaining({ + metadata: expect.objectContaining({ + models: [{ providerID: "fireworks", modelID: longModelID, variant: "default" }], + }), + }), + ); + }); }); describe("handleModelSearchResults", () => { @@ -226,5 +347,36 @@ describe("bot model selection", () => { expect(result).toBe(false); }); + + it("resolves short search result callback data to the original long model ID", async () => { + const longModelID = "accounts/hubabuba3227-1hvtqlh/deployments/kpwpvuky"; + mocked.interactionManagerGetSnapshotMock.mockReturnValue({ + kind: "custom", + metadata: { + flow: "model-search", + stage: "results", + messageId: 999, + models: [{ providerID: "fireworks", modelID: longModelID, variant: "default" }], + }, + }); + + const ctx = mockContext({ + callbackQuery: { + data: "model:result:0", + message: { message_id: 999 }, + }, + api: {}, + }); + + const result = await handleModelSearchResults(ctx); + + expect(result).toBe(true); + expect(mocked.interactionManagerClearMock).toHaveBeenCalledWith("model_search_selected"); + expect(mocked.selectModelMock).toHaveBeenCalledWith({ + providerID: "fireworks", + modelID: longModelID, + variant: "default", + }); + }); }); }); From 47214d778e8c0930602ae502539abeaf4c662d3c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E5=9F=BA=E9=AD=81?= <1412414664@qq.com> Date: Tue, 7 Jul 2026 10:58:25 +0800 Subject: [PATCH 2/3] fix(model): reject stale short callbacks --- .../model-selection-callback-handler.ts | 24 +++++-- tests/bot/handlers/model.test.ts | 63 +++++++++++++++++++ 2 files changed, 82 insertions(+), 5 deletions(-) diff --git a/src/bot/callbacks/model-selection-callback-handler.ts b/src/bot/callbacks/model-selection-callback-handler.ts index 3903e6f4..cd1b1b25 100644 --- a/src/bot/callbacks/model-selection-callback-handler.ts +++ b/src/bot/callbacks/model-selection-callback-handler.ts @@ -146,6 +146,13 @@ function parseLegacyModelCallback(data: string): ModelInfo | null { }; } +function isShortModelCallback(data: string): boolean { + return ( + data.startsWith(MODEL_SEARCH_RESULT_CALLBACK_PREFIX) || + data.startsWith(MODEL_LIST_CALLBACK_PREFIX) + ); +} + /** * Shared logic for applying a model selection and updating UI. * Used by both the regular inline menu flow and the search results flow. @@ -217,11 +224,12 @@ export async function handleModelSelect(ctx: Context): Promise { logger.debug(`[ModelHandler] Received callback: ${callbackQuery.data}`); try { - const modelInfo = - (await resolveModelListCallback(callbackQuery.data)) ?? - parseLegacyModelCallback(callbackQuery.data); + const modelInfo = await resolveModelListCallback(callbackQuery.data); + const shouldUseLegacyFallback = !isShortModelCallback(callbackQuery.data); + const resolvedModelInfo = + modelInfo ?? (shouldUseLegacyFallback ? parseLegacyModelCallback(callbackQuery.data) : null); - if (!modelInfo) { + if (!resolvedModelInfo) { logger.error(`[ModelHandler] Invalid callback data format: ${callbackQuery.data}`); clearActiveInlineMenu("model_select_invalid_callback"); await ctx.answerCallbackQuery({ text: t("model.change_error_callback") }).catch(() => {}); @@ -229,7 +237,7 @@ export async function handleModelSelect(ctx: Context): Promise { } clearActiveInlineMenu("model_selected"); - await applyModelSelectionAndNotify(ctx, modelInfo); + await applyModelSelectionAndNotify(ctx, resolvedModelInfo); return true; } catch (err) { @@ -410,6 +418,12 @@ export async function handleModelSearchResults(ctx: Context): Promise { // Backward compatibility for callbacks from already-rendered search result messages. if (data.startsWith("model:")) { + if (isShortModelCallback(data)) { + logger.error(`[ModelHandler] Invalid search result callback data: ${data}`); + await ctx.answerCallbackQuery({ text: t("model.change_error_callback") }).catch(() => {}); + return true; + } + const modelInfo = parseLegacyModelCallback(data); if (!modelInfo) { return true; diff --git a/tests/bot/handlers/model.test.ts b/tests/bot/handlers/model.test.ts index 563d1e92..533f0415 100644 --- a/tests/bot/handlers/model.test.ts +++ b/tests/bot/handlers/model.test.ts @@ -183,6 +183,43 @@ describe("bot model selection", () => { variant: "default", }); }); + + it("rejects stale search result callbacks instead of parsing them as legacy models", async () => { + const ctx = mockContext({ + callbackQuery: { + data: "model:result:0", + message: { message_id: 999 }, + }, + api: {}, + }); + + const result = await handleModelSelect(ctx); + + expect(result).toBe(true); + expect(mocked.selectModelMock).not.toHaveBeenCalled(); + expect(ctx.answerCallbackQuery).toHaveBeenCalled(); + }); + + it("rejects unresolved short list callbacks instead of parsing them as legacy models", async () => { + mocked.getModelSelectionListsMock.mockResolvedValue({ + favorites: [], + recent: [], + }); + + const ctx = mockContext({ + callbackQuery: { + data: "model:list:recent:0", + message: { message_id: 999 }, + }, + api: {}, + }); + + const result = await handleModelSelect(ctx); + + expect(result).toBe(true); + expect(mocked.selectModelMock).not.toHaveBeenCalled(); + expect(ctx.answerCallbackQuery).toHaveBeenCalled(); + }); }); describe("handleModelSearchCallback", () => { @@ -378,5 +415,31 @@ describe("bot model selection", () => { variant: "default", }); }); + + it("rejects stale short list callbacks instead of parsing them as legacy models", async () => { + mocked.interactionManagerGetSnapshotMock.mockReturnValue({ + kind: "custom", + metadata: { + flow: "model-search", + stage: "results", + messageId: 999, + models: [], + }, + }); + + const ctx = mockContext({ + callbackQuery: { + data: "model:list:recent:0", + message: { message_id: 999 }, + }, + api: {}, + }); + + const result = await handleModelSearchResults(ctx); + + expect(result).toBe(true); + expect(mocked.selectModelMock).not.toHaveBeenCalled(); + expect(ctx.answerCallbackQuery).toHaveBeenCalled(); + }); }); }); From f05fd0494621e4d67b42ae23e672c459635312f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E5=9F=BA=E9=AD=81?= <1412414664@qq.com> Date: Fri, 10 Jul 2026 22:16:54 +0800 Subject: [PATCH 3/3] fix(model): bind list callbacks to menu snapshot --- .../model-selection-callback-handler.ts | 38 ++++++++++--- src/bot/menus/inline-menu.ts | 4 +- src/bot/menus/model-selection-menu.ts | 1 + tests/bot/handlers/model.test.ts | 54 ++++++++++++++++--- tests/bot/menus/inline-menu.test.ts | 10 ++++ 5 files changed, 91 insertions(+), 16 deletions(-) diff --git a/src/bot/callbacks/model-selection-callback-handler.ts b/src/bot/callbacks/model-selection-callback-handler.ts index cd1b1b25..a900ca8e 100644 --- a/src/bot/callbacks/model-selection-callback-handler.ts +++ b/src/bot/callbacks/model-selection-callback-handler.ts @@ -1,10 +1,6 @@ import { Context, InlineKeyboard } from "grammy"; import { getStoredAgent, resolveProjectAgent } from "../../app/services/agent-selection-service.js"; -import { - getModelSelectionLists, - searchModels, - selectModel, -} from "../../app/services/model-selection-service.js"; +import { searchModels, selectModel } from "../../app/services/model-selection-service.js"; import { formatVariantForButton } from "../../app/services/variant-selection-service.js"; import { formatModelForDisplay } from "../../app/types/model.js"; import type { ModelInfo } from "../../app/types/model.js"; @@ -31,6 +27,11 @@ interface ModelSearchMetadata { models: ModelInfo[]; } +interface ModelListMetadata { + favorites: ModelInfo[]; + recent: ModelInfo[]; +} + function parseModelItems(value: unknown): ModelInfo[] { if (!Array.isArray(value)) { return []; @@ -77,6 +78,23 @@ function parseModelSearchMetadata(): ModelSearchMetadata | null { return { flow, stage, messageId, models: parseModelItems(state.metadata.models) }; } +function parseModelListMetadata(): ModelListMetadata | null { + const state = interactionManager.getSnapshot(); + if (!state || state.kind !== "inline" || state.metadata.menuKind !== "model") { + return null; + } + + const modelLists = state.metadata.modelLists; + if (typeof modelLists !== "object" || modelLists === null) { + return null; + } + + return { + favorites: parseModelItems("favorites" in modelLists ? modelLists.favorites : undefined), + recent: parseModelItems("recent" in modelLists ? modelLists.recent : undefined), + }; +} + function parseNonNegativeIndex(value: string): number | null { if (!/^\d+$/.test(value)) { return null; @@ -98,7 +116,7 @@ function parseCallbackIndex(data: string, prefix: string): number | null { return parseNonNegativeIndex(data.slice(prefix.length)); } -async function resolveModelListCallback(data: string): Promise { +function resolveModelListCallback(data: string): ModelInfo | null { if (!data.startsWith(MODEL_LIST_CALLBACK_PREFIX)) { return null; } @@ -114,7 +132,11 @@ async function resolveModelListCallback(data: string): Promise return null; } - const lists = await getModelSelectionLists(); + const lists = parseModelListMetadata(); + if (!lists) { + return null; + } + const model = kind === "favorites" ? lists.favorites[index] : lists.recent[index]; if (!model) { return null; @@ -224,7 +246,7 @@ export async function handleModelSelect(ctx: Context): Promise { logger.debug(`[ModelHandler] Received callback: ${callbackQuery.data}`); try { - const modelInfo = await resolveModelListCallback(callbackQuery.data); + const modelInfo = resolveModelListCallback(callbackQuery.data); const shouldUseLegacyFallback = !isShortModelCallback(callbackQuery.data); const resolvedModelInfo = modelInfo ?? (shouldUseLegacyFallback ? parseLegacyModelCallback(callbackQuery.data) : null); diff --git a/src/bot/menus/inline-menu.ts b/src/bot/menus/inline-menu.ts index 18f7984e..f17c7e7f 100644 --- a/src/bot/menus/inline-menu.ts +++ b/src/bot/menus/inline-menu.ts @@ -1,6 +1,6 @@ import { Context, InlineKeyboard } from "grammy"; import { interactionManager } from "../../app/managers/interaction-manager.js"; -import type { InteractionState } from "../../app/types/interaction.js"; +import type { InteractionMetadata, InteractionState } from "../../app/types/interaction.js"; import { logger } from "../../utils/logger.js"; import { t } from "../../i18n/index.js"; @@ -32,6 +32,7 @@ interface InlineMenuReplyOptions { text: string; keyboard: InlineKeyboard; parseMode?: "Markdown" | "HTML"; + metadata?: InteractionMetadata; } export function isInlineMenuKind(value: string): value is InlineMenuKind { @@ -118,6 +119,7 @@ export async function replyWithInlineMenu( kind: "inline", expectedInput: "callback", metadata: { + ...options.metadata, menuKind: options.menuKind, messageId: message.message_id, }, diff --git a/src/bot/menus/model-selection-menu.ts b/src/bot/menus/model-selection-menu.ts index 3c0b969c..14c5274c 100644 --- a/src/bot/menus/model-selection-menu.ts +++ b/src/bot/menus/model-selection-menu.ts @@ -94,6 +94,7 @@ export async function showModelSelectionMenu(ctx: Context): Promise { menuKind: "model", text, keyboard, + metadata: { modelLists }, }); } catch (err) { logger.error("[ModelHandler] Error showing model menu:", err); diff --git a/tests/bot/handlers/model.test.ts b/tests/bot/handlers/model.test.ts index 533f0415..21949144 100644 --- a/tests/bot/handlers/model.test.ts +++ b/tests/bot/handlers/model.test.ts @@ -19,6 +19,7 @@ const mocked = vi.hoisted(() => ({ pinnedGetContextInfoMock: vi.fn(), pinnedGetContextLimitMock: vi.fn(), createMainKeyboardMock: vi.fn(), + replyWithInlineMenuMock: vi.fn(), })); vi.mock("../../../src/app/services/model-selection-service.js", () => ({ @@ -66,10 +67,13 @@ vi.mock("../../../src/app/managers/interaction-manager.js", () => ({ vi.mock("../../../src/bot/menus/inline-menu.js", () => ({ ensureActiveInlineMenu: mocked.ensureActiveInlineMenuMock, clearActiveInlineMenu: vi.fn(), - replyWithInlineMenu: vi.fn(), + replyWithInlineMenu: mocked.replyWithInlineMenuMock, })); -import { buildModelSelectionMenu } from "../../../src/bot/menus/model-selection-menu.js"; +import { + buildModelSelectionMenu, + showModelSelectionMenu, +} from "../../../src/bot/menus/model-selection-menu.js"; import { handleModelSelect, @@ -110,6 +114,7 @@ describe("bot model selection", () => { mocked.pinnedGetContextInfoMock.mockReset().mockReturnValue(null); mocked.pinnedGetContextLimitMock.mockReset().mockReturnValue(0); mocked.createMainKeyboardMock.mockReset().mockReturnValue({ keyboard: [["main"]] }); + mocked.replyWithInlineMenuMock.mockReset().mockResolvedValue(999); }); describe("buildModelSelectionMenu", () => { @@ -156,14 +161,44 @@ describe("bot model selection", () => { expect(Buffer.byteLength(callbackData ?? "", "utf-8")).toBeLessThanOrEqual(64); expect(callbackData).not.toContain(longModelID); }); + + it("stores the rendered model lists with the active menu", async () => { + const modelLists = { + favorites: [{ providerID: "openai", modelID: "gpt-4o" }], + recent: [{ providerID: "google", modelID: "gemini-pro" }], + }; + mocked.getModelSelectionListsMock.mockResolvedValue(modelLists); + const ctx = mockContext(); + + await showModelSelectionMenu(ctx); + + expect(mocked.replyWithInlineMenuMock).toHaveBeenCalledWith( + ctx, + expect.objectContaining({ + menuKind: "model", + metadata: { modelLists }, + }), + ); + }); }); describe("handleModelSelect", () => { - it("resolves short list callback data to the original long model ID", async () => { + it("resolves short list callback data from the rendered menu snapshot", async () => { const longModelID = "accounts/hubabuba3227-1hvtqlh/deployments/kpwpvuky"; + mocked.interactionManagerGetSnapshotMock.mockReturnValue({ + kind: "inline", + metadata: { + menuKind: "model", + messageId: 999, + modelLists: { + favorites: [], + recent: [{ providerID: "fireworks", modelID: longModelID }], + }, + }, + }); mocked.getModelSelectionListsMock.mockResolvedValue({ favorites: [], - recent: [{ providerID: "fireworks", modelID: longModelID }], + recent: [{ providerID: "openai", modelID: "different-model" }], }); const ctx = mockContext({ @@ -182,6 +217,7 @@ describe("bot model selection", () => { modelID: longModelID, variant: "default", }); + expect(mocked.getModelSelectionListsMock).not.toHaveBeenCalled(); }); it("rejects stale search result callbacks instead of parsing them as legacy models", async () => { @@ -201,9 +237,13 @@ describe("bot model selection", () => { }); it("rejects unresolved short list callbacks instead of parsing them as legacy models", async () => { - mocked.getModelSelectionListsMock.mockResolvedValue({ - favorites: [], - recent: [], + mocked.interactionManagerGetSnapshotMock.mockReturnValue({ + kind: "inline", + metadata: { + menuKind: "model", + messageId: 999, + modelLists: { favorites: [], recent: [] }, + }, }); const ctx = mockContext({ diff --git a/tests/bot/menus/inline-menu.test.ts b/tests/bot/menus/inline-menu.test.ts index 837973a8..1601d6cb 100644 --- a/tests/bot/menus/inline-menu.test.ts +++ b/tests/bot/menus/inline-menu.test.ts @@ -87,6 +87,12 @@ describe("bot/menus/inline-menu", () => { menuKind: "model", text: "Select model", keyboard, + metadata: { + modelLists: { + favorites: [{ providerID: "openai", modelID: "gpt-4o" }], + recent: [], + }, + }, }); expect(ctx.reply).toHaveBeenCalledTimes(1); @@ -102,6 +108,10 @@ describe("bot/menus/inline-menu", () => { expect(state?.expectedInput).toBe("callback"); expect(state?.metadata.menuKind).toBe("model"); expect(state?.metadata.messageId).toBe(42); + expect(state?.metadata.modelLists).toEqual({ + favorites: [{ providerID: "openai", modelID: "gpt-4o" }], + recent: [], + }); }); it("accepts callback from active inline menu", async () => {