diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts index 971c8ada5d..79eb8325b1 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts @@ -119,6 +119,8 @@ export type AgentMode = string; type PrepareSessionSharedFields = { mode: AgentMode; model: string; + /** Optional cheap same-vendor model for title/aux calls (Code Reviewer). */ + smallModel?: string; variant?: string; // GitHub-specific params githubRepo?: string; diff --git a/apps/web/src/lib/code-reviews/core/model-selection.test.ts b/apps/web/src/lib/code-reviews/core/model-selection.test.ts index f6cc0d36a3..8179c3e02d 100644 --- a/apps/web/src/lib/code-reviews/core/model-selection.test.ts +++ b/apps/web/src/lib/code-reviews/core/model-selection.test.ts @@ -1,5 +1,10 @@ -import { resolveEffectiveModel, selectedModelFromReviewSources } from './model-selection'; -import type { CodeReviewAgentConfig } from '@kilocode/db/schema-types'; +import { + catalogPricesFromStoredModels, + resolveCheapSameVendorSmallModel, + resolveEffectiveModel, + selectedModelFromReviewSources, +} from './model-selection'; +import type { CodeReviewAgentConfig, StoredModel } from '@kilocode/db/schema-types'; const FALLBACK = 'anthropic/claude-sonnet-4.6'; @@ -15,6 +20,15 @@ function baseConfig( }; } +function priced(id: string, prompt: string): StoredModel { + return { + id, + name: id, + type: 'language', + endpoints: [{ pricing: { prompt, completion: prompt } }], + }; +} + describe('resolveEffectiveModel', () => { it('uses the global model when there are no overrides', () => { const result = resolveEffectiveModel(baseConfig(), 'acme/api', FALLBACK); @@ -153,3 +167,52 @@ describe('selectedModelFromReviewSources', () => { ).toBeNull(); }); }); + +describe('resolveCheapSameVendorSmallModel', () => { + const catalog = catalogPricesFromStoredModels({ + 'anthropic/claude-sonnet-4.6': priced('anthropic/claude-sonnet-4.6', '0.000003'), + 'anthropic/claude-haiku-4.5': priced('anthropic/claude-haiku-4.5', '0.0000008'), + 'anthropic/claude-opus-4.6': priced('anthropic/claude-opus-4.6', '0.000015'), + 'openai/gpt-5': priced('openai/gpt-5', '0.00001'), + 'openai/gpt-5-nano': priced('openai/gpt-5-nano', '0.0000001'), + }); + + it('picks the cheapest strictly-cheaper same-vendor sibling', () => { + expect(resolveCheapSameVendorSmallModel('anthropic/claude-sonnet-4.6', catalog)).toBe( + 'anthropic/claude-haiku-4.5' + ); + }); + + it('does not cross vendors', () => { + expect(resolveCheapSameVendorSmallModel('openai/gpt-5', catalog)).toBe('openai/gpt-5-nano'); + }); + + it('leaves managed models unset when no cheaper sibling exists', () => { + expect(resolveCheapSameVendorSmallModel('anthropic/claude-haiku-4.5', catalog)).toBeUndefined(); + }); + + it('falls back to the primary for sole-model direct-BYOK vendors', () => { + expect(resolveCheapSameVendorSmallModel('neuralwatt/glm-5.2-short', catalog)).toBe( + 'neuralwatt/glm-5.2-short' + ); + }); + + it('does not substitute a Kilo-billed catalog sibling for an uncatalogued BYOK primary', () => { + const mixedCatalog = catalogPricesFromStoredModels({ + 'neuralwatt/glm-tiny': priced('neuralwatt/glm-tiny', '0.0000001'), + }); + expect(resolveCheapSameVendorSmallModel('neuralwatt/glm-5.2-short', mixedCatalog)).toBe( + 'neuralwatt/glm-5.2-short' + ); + }); + + it('picks a cheaper BYOK sibling when catalog prices exist', () => { + const byokCatalog = catalogPricesFromStoredModels({ + 'neuralwatt/glm-5.2-short': priced('neuralwatt/glm-5.2-short', '0.000002'), + 'neuralwatt/glm-tiny': priced('neuralwatt/glm-tiny', '0.0000001'), + }); + expect(resolveCheapSameVendorSmallModel('neuralwatt/glm-5.2-short', byokCatalog)).toBe( + 'neuralwatt/glm-tiny' + ); + }); +}); diff --git a/apps/web/src/lib/code-reviews/core/model-selection.ts b/apps/web/src/lib/code-reviews/core/model-selection.ts index 503804e503..6fcff65b64 100644 --- a/apps/web/src/lib/code-reviews/core/model-selection.ts +++ b/apps/web/src/lib/code-reviews/core/model-selection.ts @@ -12,8 +12,10 @@ * or Bitbucket). See `RepositoryModelOverrideSchema`. */ -import type { CodeReviewAgentConfig } from '@kilocode/db/schema-types'; +import type { CodeReviewAgentConfig, StoredModel } from '@kilocode/db/schema-types'; import { DEFAULT_CODE_REVIEW_MODEL } from './constants'; +import { DirectUserByokInferenceProviderIdSchema } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id'; +import { getOpenRouterModelsMetadataFromDatabase } from '@/lib/ai-gateway/providers/gateway-models-cache'; export type EffectiveModelSelection = { modelSlug: string; @@ -21,6 +23,113 @@ export type EffectiveModelSelection = { source: 'repository_override' | 'global'; }; +/** Catalog entry used to pick a cheap same-vendor small model. */ +export type CatalogModelPrice = { + id: string; + /** USD per input token; null when the catalog has no usable prompt price. */ + promptPrice: number | null; +}; + +const DIRECT_BYOK_VENDORS = new Set(DirectUserByokInferenceProviderIdSchema.options); + +export function modelVendorId(modelId: string): string | undefined { + const vendor = modelId.split('/')[0]?.trim(); + return vendor || undefined; +} + +export function isDirectByokVendor(vendor: string | undefined): boolean { + return vendor != null && DIRECT_BYOK_VENDORS.has(vendor); +} + +/** + * Build catalog price rows from the OpenRouter `StoredModel` map. + * Uses the cheapest endpoint prompt price per model. + */ +export function catalogPricesFromStoredModels( + models: Record +): CatalogModelPrice[] { + return Object.values(models) + .filter(model => (model.type ?? 'language') === 'language' && model.endpoints.length > 0) + .map(model => { + const prices = model.endpoints + .map(endpoint => + endpoint.pricing?.prompt != null ? Number.parseFloat(endpoint.pricing.prompt) : Number.NaN + ) + .filter(price => Number.isFinite(price) && price >= 0); + return { + id: model.id, + promptPrice: prices.length > 0 ? Math.min(...prices) : null, + }; + }); +} + +/** + * Pick a cheap same-vendor model for Code Reviewer title/aux calls. + * + * - Prefer the lowest-priced same-vendor sibling that is strictly cheaper than the + * primary when the primary's price is known. + * - When the primary has no catalog price, pick the cheapest same-vendor sibling; + * direct-BYOK primaries are not in the Kilo-billed catalog, so they fall back to + * the primary instead of a same-prefix catalog entry. + * - Direct-BYOK vendors with no cheaper sibling fall back to the primary (user's + * key) so aux calls do not fall through to kilo-auto/small → Gemma on Kilo credits. + * - Managed vendors with no cheaper sibling leave small_model unset. + */ +export function resolveCheapSameVendorSmallModel( + primaryModelId: string, + catalog: readonly CatalogModelPrice[] +): string | undefined { + const vendor = modelVendorId(primaryModelId); + if (!vendor) return undefined; + + const byok = isDirectByokVendor(vendor); + const primaryEntry = catalog.find(entry => entry.id === primaryModelId); + const primaryPrice = primaryEntry?.promptPrice ?? undefined; + + // Direct-BYOK models are served from the user's own provider, not the + // Kilo-billed OpenRouter catalog. A same-prefix catalog entry therefore does + // not represent the same provider, so never substitute one for a BYOK primary. + const siblings = + byok && primaryEntry == null + ? [] + : catalog.filter( + entry => + entry.id !== primaryModelId && + modelVendorId(entry.id) === vendor && + entry.promptPrice != null && + Number.isFinite(entry.promptPrice) + ); + + const cheaper = + primaryPrice != null && Number.isFinite(primaryPrice) + ? siblings.filter(entry => (entry.promptPrice as number) < primaryPrice) + : siblings; + + if (cheaper.length === 0) { + return byok ? primaryModelId : undefined; + } + + cheaper.sort((a, b) => { + const priceDelta = (a.promptPrice as number) - (b.promptPrice as number); + return priceDelta !== 0 ? priceDelta : a.id.localeCompare(b.id); + }); + return cheaper[0]?.id; +} + +/** + * Resolve the Code Reviewer aux/title small model for a primary review model. + * Soft-fails to BYOK-primary or unset when the gateway catalog cannot be loaded. + */ +export async function resolveReviewSmallModel(primaryModelId: string): Promise { + const vendor = modelVendorId(primaryModelId); + try { + const models = await getOpenRouterModelsMetadataFromDatabase(); + return resolveCheapSameVendorSmallModel(primaryModelId, catalogPricesFromStoredModels(models)); + } catch { + return isDirectByokVendor(vendor) ? primaryModelId : undefined; + } +} + /** * Resolve the effective model for a review's repository. * diff --git a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.test.ts b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.test.ts index 2f85c05e2b..d0db527bd2 100644 --- a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.test.ts +++ b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.test.ts @@ -15,9 +15,11 @@ const mockFindPreviousCompletedReview = jest.fn(); const mockUpdatePreviousReviewSummary = jest.fn(); const mockUpdateRepositoryReviewInstructionsMetadata = jest.fn(); const mockGenerateReviewPrompt = jest.fn(); +const mockResolveReviewSmallModel = jest.fn(); import type { CodeReviewAgentConfig } from '@/lib/agent-config/core/types'; import type * as CodeReviewsDb from '@/lib/code-reviews/db/code-reviews'; +import type * as ModelSelection from '@/lib/code-reviews/core/model-selection'; jest.mock('@/lib/integrations/platforms/github/adapter', () => ({ generateGitHubInstallationToken: (...args: unknown[]) => @@ -62,6 +64,16 @@ jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn(), })); +jest.mock('@/lib/code-reviews/core/model-selection', () => { + const actual = jest.requireActual( + '@/lib/code-reviews/core/model-selection' + ); + return { + ...actual, + resolveReviewSmallModel: (...args: unknown[]) => mockResolveReviewSmallModel(...args), + }; +}); + import { db } from '@/lib/drizzle'; import { insertTestUser } from '@/tests/helpers/user.helper'; import { @@ -249,6 +261,8 @@ describe('prepareReviewPayload', () => { mockUpdatePreviousReviewSummary.mockReset(); mockUpdateRepositoryReviewInstructionsMetadata.mockReset(); mockGenerateReviewPrompt.mockReset(); + mockResolveReviewSmallModel.mockReset(); + mockResolveReviewSmallModel.mockResolvedValue(undefined); }); afterAll(async () => { @@ -979,4 +993,39 @@ describe('prepareReviewPayload', () => { upstreamBranch: 'refs/pull/1235/head', }); }); + + it('includes resolved smallModel on GitHub session input when present', async () => { + const [review] = await db + .insert(cloud_agent_code_reviews) + .values(defineReview(testUser.id, integration.id)) + .returning(); + mockResolveReviewSmallModel.mockResolvedValueOnce('anthropic/claude-haiku-4.5'); + + const payload = await prepareReviewPayload({ + reviewId: review.id, + owner: { type: 'user', id: testUser.id, userId: testUser.id }, + agentConfig: { config: baseAgentConfig }, + platform: 'github', + }); + + expect(mockResolveReviewSmallModel).toHaveBeenCalledWith('test-model'); + expect(payload.sessionInput.smallModel).toBe('anthropic/claude-haiku-4.5'); + }); + + it('omits smallModel from session input when no cheap sibling is resolved', async () => { + const [review] = await db + .insert(cloud_agent_code_reviews) + .values(defineReview(testUser.id, integration.id)) + .returning(); + + const payload = await prepareReviewPayload({ + reviewId: review.id, + owner: { type: 'user', id: testUser.id, userId: testUser.id }, + agentConfig: { config: baseAgentConfig }, + platform: 'github', + }); + + expect(mockResolveReviewSmallModel).toHaveBeenCalledWith('test-model'); + expect(payload.sessionInput).not.toHaveProperty('smallModel'); + }); }); diff --git a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts index 65ee241b82..4a36b53eeb 100644 --- a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts +++ b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts @@ -60,6 +60,7 @@ import { type ReviewScope, } from '../db/code-reviews'; import { DEFAULT_CODE_REVIEW_MODEL, DEFAULT_CODE_REVIEW_MODE } from '../core/constants'; +import { resolveReviewSmallModel } from '../core/model-selection'; import type { Owner } from '../core'; import { generateReviewPrompt } from '../prompts/generate-prompt'; import type { CodeReviewAgentConfig } from '@/lib/agent-config/core/types'; @@ -113,6 +114,11 @@ export type SessionInput = { prompt: string; mode: 'code'; model: string; + /** + * Optional cheap same-vendor model for kilo title/aux calls. When omitted, the + * CLI falls through to its default small-model list (kilo-auto/small → Gemma). + */ + smallModel?: string; /** Thinking effort variant name (e.g. "high", "max") — undefined means model default */ variant?: string; upstreamBranch: string; @@ -302,12 +308,14 @@ export async function prepareReviewPayload( // Single source for the standard reviewer's model so the session input and the // forward-shaped `reviewAgents[0]` can never drift apart. const standardModel = config.model_slug || DEFAULT_CODE_REVIEW_MODEL; + const smallModel = await resolveReviewSmallModel(standardModel); const sessionInput: SessionInput = { gitUrl: `https://bitbucket.org/${workspaceSlug.data}/${repositorySlug.data}.git`, kilocodeOrganizationId: owner.id, prompt, mode: DEFAULT_CODE_REVIEW_MODE as 'code', model: standardModel, + ...(smallModel ? { smallModel } : {}), variant: config.thinking_effort ?? undefined, upstreamBranch: review.head_ref, platform: PLATFORM.BITBUCKET, @@ -758,6 +766,7 @@ export async function prepareReviewPayload( // Single source for the standard reviewer's model so the session input and the // forward-shaped `reviewAgents[0]` can never drift apart. const standardModel = config.model_slug || DEFAULT_CODE_REVIEW_MODEL; + const smallModel = await resolveReviewSmallModel(standardModel); const gateThreshold = config.gate_threshold ?? 'off'; const githubCheckoutRef = getGitHubPullRequestCheckoutRef(review.pr_number); const sessionInput: SessionInput = @@ -771,6 +780,7 @@ export async function prepareReviewPayload( prompt, mode: DEFAULT_CODE_REVIEW_MODE as 'code', model: standardModel, + ...(smallModel ? { smallModel } : {}), variant, upstreamBranch: review.head_ref, } @@ -784,6 +794,7 @@ export async function prepareReviewPayload( prompt, mode: DEFAULT_CODE_REVIEW_MODE as 'code', model: standardModel, + ...(smallModel ? { smallModel } : {}), variant, upstreamBranch: review.head_ref, ...(gateThreshold !== 'off' ? { gateThreshold } : {}), @@ -797,6 +808,7 @@ export async function prepareReviewPayload( prompt, mode: DEFAULT_CODE_REVIEW_MODE as 'code', model: standardModel, + ...(smallModel ? { smallModel } : {}), variant, upstreamBranch: githubCheckoutRef, ...(gateThreshold !== 'off' ? { gateThreshold } : {}), diff --git a/packages/worker-utils/src/cloud-agent-next-client.ts b/packages/worker-utils/src/cloud-agent-next-client.ts index 3d15f59722..d2890aceeb 100644 --- a/packages/worker-utils/src/cloud-agent-next-client.ts +++ b/packages/worker-utils/src/cloud-agent-next-client.ts @@ -44,6 +44,8 @@ export type CloudAgentPrepareSessionInput = { prompt: string; mode: string; model: string; + /** Optional cheap same-vendor model for title/aux calls (Code Reviewer). */ + smallModel?: string; variant?: string; githubRepo?: string; githubToken?: string; diff --git a/services/cloud-agent-next/src/execution/types.ts b/services/cloud-agent-next/src/execution/types.ts index 8fd5fc64ca..fe146e1005 100644 --- a/services/cloud-agent-next/src/execution/types.ts +++ b/services/cloud-agent-next/src/execution/types.ts @@ -85,6 +85,8 @@ export function renderExecutionTurnContent(turn: AcceptedExecutionTurn): string export type ModelChoice = { model: string; variant?: string; + /** Optional cheap same-vendor model for title/aux calls (Code Reviewer). */ + smallModel?: string; }; /** Fully resolved agent selection. */ diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index b567e6aae9..16b66d701c 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -447,6 +447,7 @@ function isSameInitialAdmissionConfiguration( JSON.stringify(input.workspace?.sandboxRoute) && metadata.agent?.mode === input.agent.mode && metadata.agent.model === input.agent.model && + metadata.agent.smallModel === input.agent.smallModel && metadata.agent.variant === input.agent.variant && metadata.finalization?.autoCommit === input.finalization?.autoCommit && metadata.finalization?.condenseOnComplete === input.finalization?.condenseOnComplete @@ -2551,6 +2552,7 @@ export class CloudAgentSession extends DurableObject { agent: { mode: input.agent.mode, model: input.agent.model, + smallModel: input.agent.smallModel, variant: input.agent.variant, appendSystemPrompt: input.agent.appendSystemPrompt, }, @@ -3766,6 +3768,7 @@ export class CloudAgentSession extends DurableObject { agent: { mode: metadata.agent.mode, model: metadata.agent.model, + smallModel: metadata.agent.smallModel, variant: metadata.agent.variant, }, finalization: { diff --git a/services/cloud-agent-next/src/persistence/session-metadata.test.ts b/services/cloud-agent-next/src/persistence/session-metadata.test.ts index e9c5036a3b..af81c61fec 100644 --- a/services/cloud-agent-next/src/persistence/session-metadata.test.ts +++ b/services/cloud-agent-next/src/persistence/session-metadata.test.ts @@ -272,6 +272,7 @@ describe('session metadata boundary', () => { agent: { mode: 'reviewer', model: 'kilo/gpt-5', + smallModel: 'kilo/anthropic/claude-haiku-4.5', variant: 'thinking', appendSystemPrompt: 'Extra context', }, diff --git a/services/cloud-agent-next/src/persistence/session-metadata.ts b/services/cloud-agent-next/src/persistence/session-metadata.ts index 63f1f55353..b14838b0eb 100644 --- a/services/cloud-agent-next/src/persistence/session-metadata.ts +++ b/services/cloud-agent-next/src/persistence/session-metadata.ts @@ -193,6 +193,7 @@ const MetadataAgentSchema = z .object({ mode: z.string().optional(), model: z.string().optional(), + smallModel: z.string().optional(), variant: z .string() .max(50) diff --git a/services/cloud-agent-next/src/router/handlers/session-prepare.ts b/services/cloud-agent-next/src/router/handlers/session-prepare.ts index 007ab40850..6fa6ae18fe 100644 --- a/services/cloud-agent-next/src/router/handlers/session-prepare.ts +++ b/services/cloud-agent-next/src/router/handlers/session-prepare.ts @@ -270,6 +270,7 @@ export function prepareInputToSessionCreateRequest(input: PrepareInput): Session agent: { mode: input.mode, model: input.model, + ...(input.smallModel ? { smallModel: input.smallModel } : {}), variant: input.variant, }, repository, diff --git a/services/cloud-agent-next/src/router/schemas.ts b/services/cloud-agent-next/src/router/schemas.ts index 35eb78c7db..61bd23d7bc 100644 --- a/services/cloud-agent-next/src/router/schemas.ts +++ b/services/cloud-agent-next/src/router/schemas.ts @@ -436,6 +436,11 @@ const PrepareSessionSharedFields = { 'Kilo Code execution mode (built-in or custom slug from runtimeAgents)' ), model: modelIdSchema.describe('AI model to use'), + smallModel: modelIdSchema + .optional() + .describe( + 'Optional cheap same-vendor model for title/aux calls (Code Reviewer). When omitted, CLI defaults apply.' + ), variant: z .string() .max(50) diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index dd988ba3da..ff2fd82bb6 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -178,6 +178,32 @@ describe('SessionService.buildRuntimeEnv', () => { expect(runtimeEnv.GIT_CONFIG_GLOBAL).toBeUndefined(); expect(runtimeEnv.GIT_CONFIG_NOSYSTEM).toBeUndefined(); }); + + it('materializes code-review small_model from the smallModel option', () => { + const service = new SessionService(); + const context = service.buildContext({ + sandboxId: 'usr-test', + userId: 'user_test', + sessionId: 'agent_test', + }); + const smallModel = 'anthropic/claude-haiku-4.5'; + + const runtimeEnv = service.buildRuntimeEnv({ + context, + env: createEnv(), + kiloCapability: 'kilo-token', + kilocodeModel: 'anthropic/claude-sonnet-4.6', + smallModel, + createdOnPlatform: 'code-review', + }); + + const config = JSON.parse(runtimeEnv.KILO_CONFIG_CONTENT) as { + small_model?: string; + agent?: { title?: { model?: string } }; + }; + expect(config.small_model).toBe(`kilo/${smallModel}`); + expect(config.agent?.title?.model).toBe(`kilo/${smallModel}`); + }); }); describe('code-review command guard policy', () => { @@ -2444,6 +2470,140 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { }); } + it('pins code-review KILO_CONFIG small_model to the dispatched smallModel', async () => { + const service = new SessionService(); + const env = createEnv(); + env.WORKER_URL = 'https://cloud-agent.example.com'; + const primaryModel = 'anthropic/claude-sonnet-4.6'; + const smallModel = 'anthropic/claude-haiku-4.5'; + + const build = (args: { + messageId: string; + wrapperRunId: string; + metadata: CloudAgentSessionState; + agent: { mode: 'code'; model: string; smallModel?: string }; + }) => + service.buildWrapperSessionReadyAndPromptRequests({ + env, + plan: { + scope: { sessionId: 'agent_test', userId: 'user_test' }, + turn: { type: 'prompt', messageId: args.messageId, prompt: 'Review the PR' }, + agent: args.agent, + workspace: { + sandboxId: args.metadata.workspace?.sandboxId ?? 'ses-abcdef', + metadata: args.metadata, + }, + wrapper: { + fence: { + wrapperRunId: args.wrapperRunId, + wrapperGeneration: 1, + wrapperConnectionId: `conn_${args.wrapperRunId}`, + }, + }, + } satisfies FencedWrapperDispatchRequest, + }); + + const readConfig = (result: Awaited>) => + JSON.parse(result.readyRequest.materialized.env.KILO_CONFIG_CONTENT) as { + model?: string; + small_model?: string; + agent?: { title?: { model?: string } }; + }; + + const codeReview = readConfig( + await build({ + messageId: 'msg_018f1e2d3c4bSmallModelAAAA', + wrapperRunId: 'wr_small_model', + metadata: createMetadata({ createdOnPlatform: 'code-review' }), + agent: { mode: 'code', model: primaryModel, smallModel }, + }) + ); + expect(codeReview.model).toBe(`kilo/${primaryModel}`); + expect(codeReview.small_model).toBe(`kilo/${smallModel}`); + expect(codeReview.agent?.title?.model).toBe(`kilo/${smallModel}`); + + const unset = readConfig( + await build({ + messageId: 'msg_018f1e2d3c4bSmallModelNone', + wrapperRunId: 'wr_small_model_unset', + metadata: createMetadata({ createdOnPlatform: 'code-review' }), + agent: { mode: 'code', model: primaryModel }, + }) + ); + expect(unset.model).toBe(`kilo/${primaryModel}`); + expect(unset.small_model).toBeUndefined(); + expect(unset.agent?.title).toBeUndefined(); + + const web = readConfig( + await build({ + messageId: 'msg_018f1e2d3c4bSmallModelBBBB', + wrapperRunId: 'wr_small_model_web', + metadata: createMetadata({ createdOnPlatform: 'cloud-agent-web' }), + agent: { mode: 'code', model: primaryModel, smallModel }, + }) + ); + expect(web.model).toBe(`kilo/${primaryModel}`); + expect(web.small_model).toBeUndefined(); + expect(web.agent?.title).toBeUndefined(); + + // Production code-review sessions persist smallModel in metadata; the plan + // agent does not carry it, so this exercises the metadata fallback. + const baseMetadata = createMetadata({ createdOnPlatform: 'code-review' }); + const fromMetadata = readConfig( + await build({ + messageId: 'msg_018f1e2d3c4bSmallModelCCCC', + wrapperRunId: 'wr_small_model_metadata', + metadata: { ...baseMetadata, agent: { ...baseMetadata.agent, smallModel } }, + agent: { mode: 'code', model: primaryModel }, + }) + ); + expect(fromMetadata.model).toBe(`kilo/${primaryModel}`); + expect(fromMetadata.small_model).toBe(`kilo/${smallModel}`); + expect(fromMetadata.agent?.title?.model).toBe(`kilo/${smallModel}`); + }); + + it('sets Bitbucket code-review small_model without injecting agent.title', async () => { + const service = new SessionService(); + const env = createEnv(); + env.WORKER_URL = 'https://cloud-agent.example.com'; + const primaryModel = 'anthropic/claude-sonnet-4.6'; + const smallModel = 'anthropic/claude-haiku-4.5'; + const metadata = createBitbucketMetadata(true); + + const result = await service.buildWrapperSessionReadyAndPromptRequests({ + env, + plan: { + scope: { sessionId: 'agent_test', userId: 'user_test' }, + turn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bSmallModelBBBB', + prompt: 'Review the PR', + }, + agent: { mode: 'code', model: primaryModel, smallModel }, + workspace: { + sandboxId: metadata.workspace?.sandboxId ?? 'ses-abcdef', + metadata, + }, + wrapper: { + fence: { + wrapperRunId: 'wr_bb_small_model', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_bb_small_model', + }, + }, + } satisfies FencedWrapperDispatchRequest, + }); + + const config = JSON.parse(result.readyRequest.materialized.env.KILO_CONFIG_CONTENT) as { + model?: string; + small_model?: string; + agent?: { title?: { model?: string } }; + }; + expect(config.model).toBe(`kilo/${primaryModel}`); + expect(config.small_model).toBe(`kilo/${smallModel}`); + expect(config.agent?.title).toBeUndefined(); + }); + it('uses the persisted shared checkout when constructing wrapper requests', async () => { const worktreeId = 'worktree_420ae020-e3c4-4e67-878b-66672c3d997e'; const workspacePath = `/workspace/user_test/worktrees/${worktreeId}`; diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index 3cb68c4038..bfa49efde2 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -1282,6 +1282,7 @@ export class SessionService { kiloProviderBaseUrl: opts.kiloProviderBaseUrl, kiloSessionIngestBaseUrl: opts.kiloSessionIngestBaseUrl, kilocodeModel: opts.kilocodeModel, + smallModel: opts.smallModel, originalOrgId: opts.originalOrgId, githubToken: context.githubToken, githubRepo: context.githubRepo, @@ -1310,6 +1311,7 @@ export class SessionService { kiloProviderBaseUrl, kiloSessionIngestBaseUrl, kilocodeModel, + smallModel, originalOrgId, githubToken, githubRepo, @@ -1535,6 +1537,25 @@ export class SessionService { agentCount: runtimeAgents.length, }); } + // Code Review sessions: pin small_model (and title agent when allowed) to the + // cheap same-vendor model selected at dispatch time. Leaving this unset falls + // through to kilo-auto/small → Gemma. See https://github.com/Kilo-Org/cloud/issues/4268 + const normalizedSmallModel = + smallModel && smallModel.trim() ? normalizeKilocodeModel(smallModel) : undefined; + if (createdOnPlatform === 'code-review' && normalizedSmallModel) { + configContent.small_model = normalizedSmallModel; + // Match MCP/runtimeAgents: do not inject agent.title on Bitbucket code-review + // sessions (title falls through to getSmallModel() via small_model anyway). + if (!bitbucketInputPath) { + const existingTitle = + agentConfig.title != null && + typeof agentConfig.title === 'object' && + !Array.isArray(agentConfig.title) + ? (agentConfig.title as Record) + : {}; + agentConfig.title = { ...existingTitle, model: normalizedSmallModel }; + } + } if (Object.keys(agentConfig).length > 0) { configContent.agent = agentConfig; } @@ -1669,6 +1690,7 @@ export class SessionService { kiloCapability, kiloProviderBaseUrl, kilocodeModel, + smallModel, originalOrgId, createdOnPlatform, callbackTarget, @@ -1688,6 +1710,7 @@ export class SessionService { kiloCapability, kiloProviderBaseUrl, kilocodeModel, + smallModel, originalOrgId, createdOnPlatform, callbackTarget, @@ -2121,6 +2144,7 @@ export class SessionService { kiloProviderBaseUrl, kiloSessionIngestBaseUrl, kilocodeModel: agent.model, + smallModel: metadata.agent?.smallModel ?? agent.smallModel, originalOrgId: orgId, githubToken: resolvedTokens.githubToken, githubRepo: github?.repo, @@ -2400,6 +2424,7 @@ export class SessionService { kiloProviderBaseUrl, kiloSessionIngestBaseUrl, kilocodeModel: options.kilocodeModel, + smallModel: metadata.agent?.smallModel, originalOrgId: orgId, createdOnPlatform: metadata.identity.createdOnPlatform, callbackTarget: metadata.callback?.target, @@ -2620,6 +2645,7 @@ export class SessionService { kiloProviderBaseUrl, kiloSessionIngestBaseUrl, kilocodeModel, + smallModel: metadata.agent?.smallModel, originalOrgId: orgId, createdOnPlatform: metadata.identity.createdOnPlatform, callbackTarget: metadata.callback?.target, @@ -3032,6 +3058,7 @@ export type GetOrCreateSessionOptions = { kiloProviderBaseUrl?: string; kiloSessionIngestBaseUrl?: string; kilocodeModel?: string; + smallModel?: string; originalOrgId?: string; createdOnPlatform?: string; callbackTarget?: NonNullable['target']; @@ -3060,6 +3087,8 @@ type GetSaferEnvVarsOptions = { kiloProviderBaseUrl?: string; kiloSessionIngestBaseUrl?: string; kilocodeModel?: string; + /** Optional cheap same-vendor model for Code Reviewer title/aux calls. */ + smallModel?: string; originalOrgId?: string; githubToken?: string; githubRepo?: string; diff --git a/services/cloud-agent-next/src/session/session-prepare.test.ts b/services/cloud-agent-next/src/session/session-prepare.test.ts index d6e7f8adb2..9b086e67fd 100644 --- a/services/cloud-agent-next/src/session/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session/session-prepare.test.ts @@ -4125,4 +4125,36 @@ describe('prepareInputToSessionCreateRequest clone mapping', () => { expect(request.runtime).toEqual({ sandboxAllocation: 'isolated-standard' }); }); + + it('maps a Code Reviewer smallModel into the grouped agent selection', () => { + const request = prepareInputToSessionCreateRequest({ + prompt: 'Review the PR', + mode: 'code', + model: 'anthropic/claude-sonnet-4.6', + smallModel: 'anthropic/claude-haiku-4.5', + githubRepo: 'acme/repo', + shallow: false, + devcontainer: false, + }); + + expect(request.agent).toEqual({ + mode: 'code', + model: 'anthropic/claude-sonnet-4.6', + smallModel: 'anthropic/claude-haiku-4.5', + variant: undefined, + }); + }); + + it('omits smallModel from the grouped agent selection when absent', () => { + const request = prepareInputToSessionCreateRequest({ + prompt: 'Review the PR', + mode: 'code', + model: 'anthropic/claude-sonnet-4.6', + githubRepo: 'acme/repo', + shallow: false, + devcontainer: false, + }); + + expect(request.agent).not.toHaveProperty('smallModel'); + }); }); diff --git a/services/cloud-agent-next/src/session/session-registration.ts b/services/cloud-agent-next/src/session/session-registration.ts index ec9a6671b3..c646bd32bc 100644 --- a/services/cloud-agent-next/src/session/session-registration.ts +++ b/services/cloud-agent-next/src/session/session-registration.ts @@ -1380,6 +1380,9 @@ export async function sessionCreateIntentFingerprint( mode: input.agent.mode, model: input.agent.model, variant: input.agent.variant || undefined, + // Cheap same-vendor aux model is immutable create input: it changes the + // materialized session env, so a changed value must not replay. + smallModel: input.agent.smallModel || undefined, // The DO stores the effective appended system prompt under `agent`; // it is immutable create input, so a changed system prompt must never // replay or reconcile a prior session. diff --git a/services/cloud-agent-next/test/helpers/session-setup.ts b/services/cloud-agent-next/test/helpers/session-setup.ts index 7f9821123e..ed70b6bbe0 100644 --- a/services/cloud-agent-next/test/helpers/session-setup.ts +++ b/services/cloud-agent-next/test/helpers/session-setup.ts @@ -33,6 +33,7 @@ type TestRegisterSessionInput = { prompt: string; mode: AgentMode; model: string; + smallModel?: string; variant?: string; kiloSessionId?: string; kilocodeToken?: string; @@ -147,6 +148,7 @@ export function groupedRegisterSessionInput(input: TestRegisterSessionInput): Re agent: { mode: input.mode, model: input.model, + smallModel: input.smallModel, variant: input.variant, appendSystemPrompt: input.appendSystemPrompt, }, diff --git a/services/cloud-agent-next/test/integration/session/code-review-small-model.test.ts b/services/cloud-agent-next/test/integration/session/code-review-small-model.test.ts new file mode 100644 index 0000000000..3860a0b566 --- /dev/null +++ b/services/cloud-agent-next/test/integration/session/code-review-small-model.test.ts @@ -0,0 +1,53 @@ +import { env, runInDurableObject } from 'cloudflare:test'; +import { describe, expect, it } from 'vitest'; +import { registerReadySession } from '../../helpers/session-setup.js'; + +describe('code-review small model metadata', () => { + it('persists agent.smallModel through registerSession', async () => { + const userId = 'user_small_model'; + const sessionId = 'agent_small_model'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + prompt: 'Review the PR', + mode: 'code', + model: 'anthropic/claude-sonnet-4.6', + smallModel: 'anthropic/claude-haiku-4.5', + createdOnPlatform: 'code-review', + kilocodeToken: 'token-small-model', + }); + + const metadata = await instance.getMetadata(); + expect(metadata?.agent?.smallModel).toBe('anthropic/claude-haiku-4.5'); + expect(metadata?.agent?.model).toBe('anthropic/claude-sonnet-4.6'); + }); + }); + + it('omits agent.smallModel when the caller does not provide one', async () => { + const userId = 'user_small_model_none'; + const sessionId = 'agent_small_model_none'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + prompt: 'Do the work', + mode: 'code', + model: 'anthropic/claude-sonnet-4.6', + createdOnPlatform: 'cloud-agent-web', + kilocodeToken: 'token-small-model-none', + }); + + const metadata = await instance.getMetadata(); + expect(metadata?.agent?.smallModel).toBeUndefined(); + }); + }); +}); diff --git a/services/code-review-infra/src/types.ts b/services/code-review-infra/src/types.ts index 17b42630eb..0f03444246 100644 --- a/services/code-review-infra/src/types.ts +++ b/services/code-review-infra/src/types.ts @@ -26,6 +26,8 @@ export interface SessionInput { prompt: string; mode: 'code'; model: string; + /** Optional cheap same-vendor model for kilo title/aux calls. */ + smallModel?: string; /** Thinking effort variant name (e.g. "high", "max") — undefined means model default */ variant?: string; upstreamBranch: string;