From 8feec87f4198245a486d2775e6b98d5d50fb5c05 Mon Sep 17 00:00:00 2001 From: maphew <486200+maphew@users.noreply.github.com> Date: Thu, 10 Sep 2026 02:52:40 +0000 Subject: [PATCH 1/2] fix(gastown): use the configured BYOK model for review-thread classification areThreadsBlocking() hardcoded a Workers AI Gemma call for the auto-merge review-thread classifier. Direct-BYOK towns now run the classifier on role_models.refinery ?? default_model through the Kilo gateway, so the call bills the user's own provider key. Non-BYOK towns keep the Workers AI path. A rejected BYOK call blocks without falling back to a Kilo-billed path. Refs #4268 Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com> --- .../direct-byok-provider-ids.drift.test.ts | 8 + packages/worker-utils/package.json | 3 +- .../src/direct-byok-model.test.ts | 27 ++++ .../worker-utils/src/direct-byok-model.ts | 38 +++++ services/gastown/docs/local-debug-testing.md | 8 +- .../gastown/src/dos/town/town-scm.test.ts | 152 ++++++++++++++++++ services/gastown/src/dos/town/town-scm.ts | 75 +++++++-- 7 files changed, 299 insertions(+), 12 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/providers/direct-byok/direct-byok-provider-ids.drift.test.ts create mode 100644 packages/worker-utils/src/direct-byok-model.test.ts create mode 100644 packages/worker-utils/src/direct-byok-model.ts create mode 100644 services/gastown/src/dos/town/town-scm.test.ts diff --git a/apps/web/src/lib/ai-gateway/providers/direct-byok/direct-byok-provider-ids.drift.test.ts b/apps/web/src/lib/ai-gateway/providers/direct-byok/direct-byok-provider-ids.drift.test.ts new file mode 100644 index 0000000000..fcec4ec793 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/providers/direct-byok/direct-byok-provider-ids.drift.test.ts @@ -0,0 +1,8 @@ +import { DIRECT_BYOK_PROVIDER_IDS } from '@kilocode/worker-utils/direct-byok-model'; +import { DIRECT_BYOK_PROVIDERS_META } from './direct-byok-meta'; + +it('keeps the worker-utils direct BYOK provider ids in sync with the meta list', () => { + expect([...DIRECT_BYOK_PROVIDER_IDS].sort()).toEqual( + Object.keys(DIRECT_BYOK_PROVIDERS_META).sort() + ); +}); diff --git a/packages/worker-utils/package.json b/packages/worker-utils/package.json index 4463ca04d2..7bc362d38b 100644 --- a/packages/worker-utils/package.json +++ b/packages/worker-utils/package.json @@ -39,7 +39,8 @@ "./review-agents": "./src/review-agents.ts", "./code-review-council": "./src/code-review-council.ts", "./scheduled-job-observability": "./src/scheduled-job-observability.ts", - "./r2-client": "./src/r2-client.ts" + "./r2-client": "./src/r2-client.ts", + "./direct-byok-model": "./src/direct-byok-model.ts" }, "scripts": { "test": "vitest run", diff --git a/packages/worker-utils/src/direct-byok-model.test.ts b/packages/worker-utils/src/direct-byok-model.test.ts new file mode 100644 index 0000000000..a77866fa6c --- /dev/null +++ b/packages/worker-utils/src/direct-byok-model.test.ts @@ -0,0 +1,27 @@ +import { describe, expect, it } from 'vitest'; +import { DIRECT_BYOK_PROVIDER_IDS, isDirectByokModelId } from './direct-byok-model'; + +describe('isDirectByokModelId', () => { + it('matches every listed provider id', () => { + for (const providerId of DIRECT_BYOK_PROVIDER_IDS) { + expect(isDirectByokModelId(`${providerId}/some-model`)).toBe(true); + expect(isDirectByokModelId(`${providerId.toUpperCase()}/Some-Model`)).toBe(true); + } + }); + + it('rejects non-BYOK and managed model ids', () => { + expect(isDirectByokModelId('anthropic/claude-sonnet-4.6')).toBe(false); + expect(isDirectByokModelId('google/gemma-4-26b-a4b-it')).toBe(false); + expect(isDirectByokModelId('kilo-auto/small')).toBe(false); + expect(isDirectByokModelId(undefined)).toBe(false); + expect(isDirectByokModelId(null)).toBe(false); + expect(isDirectByokModelId('')).toBe(false); + }); + + it('matches on the provider prefix only', () => { + // A bare provider id matches — routing only inspects the prefix. + expect(isDirectByokModelId('synthetic')).toBe(true); + expect(isDirectByokModelId('synthetic/hf:zai-org/GLM-5.1')).toBe(true); + expect(isDirectByokModelId('synthetic-new/whatever')).toBe(false); + }); +}); diff --git a/packages/worker-utils/src/direct-byok-model.ts b/packages/worker-utils/src/direct-byok-model.ts new file mode 100644 index 0000000000..5204716e60 --- /dev/null +++ b/packages/worker-utils/src/direct-byok-model.ts @@ -0,0 +1,38 @@ +/** + * Provider ids whose models route to the user's own API key (direct BYOK) and + * never bill Kilo credits. A model id is `/` — see + * `formatDirectByokModelId` in + * apps/web/src/lib/ai-gateway/providers/direct-byok/index.ts. + * + * Source of truth is `DIRECT_BYOK_PROVIDERS_META` in + * apps/web/src/lib/ai-gateway/providers/direct-byok/direct-byok-meta.ts. + * This copy exists because Cloudflare Workers cannot import from apps/web. + * A drift-guard test keeps the two lists equal. + */ +export const DIRECT_BYOK_PROVIDER_IDS = [ + 'alibaba-token-plan', + 'byteplus-coding', + 'chutes-byok', + 'crofai', + 'edenai', + 'kimi-coding', + 'inceptron-byok', + 'martian', + 'morph-byok', + 'neuralwatt', + 'nvidia-byok', + 'ollama-cloud', + 'opencode-go', + 'orcarouter', + 'synthetic', + 'xiaomi-token-plan-ams', + 'xiaomi-token-plan-sgp', + 'zai-coding', +] as const; + +const ids: ReadonlySet = new Set(DIRECT_BYOK_PROVIDER_IDS); + +export function isDirectByokModelId(modelId: string | undefined | null): boolean { + if (!modelId) return false; + return ids.has(modelId.toLowerCase().split('/')[0] ?? ''); +} diff --git a/services/gastown/docs/local-debug-testing.md b/services/gastown/docs/local-debug-testing.md index 3918cb1362..ecd6212a52 100644 --- a/services/gastown/docs/local-debug-testing.md +++ b/services/gastown/docs/local-debug-testing.md @@ -312,16 +312,20 @@ During testing, container restarts generate many of these. Bulk-close via admin ## 7. Auto-Merge with Workers AI Thread Classification -The auto-merge flow uses Workers AI (Gemma 4 26B) to classify unresolved PR review threads as blocking vs non-blocking. This prevents informational bot comments (status reports, code review summaries) from blocking auto-merge. +The auto-merge flow classifies unresolved PR review threads as blocking vs non-blocking. This prevents informational bot comments (status reports, code review summaries) from blocking auto-merge. ### How It Works 1. `poll_pr` runs every ~60s for MR beads with a `pr_url` 2. `checkPRFeedback` fetches review threads via GitHub GraphQL (including comment bodies) -3. If unresolved threads exist, `areThreadsBlocking()` sends them to Workers AI +3. If unresolved threads exist, `areThreadsBlocking()` classifies them: + - Direct-BYOK towns (configured model's provider prefix is a direct BYOK provider, e.g. `neuralwatt/...`) call the Kilo gateway at `/api/openrouter/chat/completions` on `role_models.refinery ?? default_model`, billing the user's own provider key + - Everyone else uses Workers AI (Gemma 4 26B) 4. The model classifies threads as BLOCKING (requires code changes, bugs, security) or NON-BLOCKING (informational, nits, bot status reports) 5. Only truly blocking threads prevent auto-merge +A rejected BYOK gateway call blocks auto-merge (`blocking=true`) rather than falling back to the Kilo-billed Workers AI path. + ### Config Required Set these on the town config (via `PATCH /debug/towns/:townId/config`): diff --git a/services/gastown/src/dos/town/town-scm.test.ts b/services/gastown/src/dos/town/town-scm.test.ts new file mode 100644 index 0000000000..fe0e90dc94 --- /dev/null +++ b/services/gastown/src/dos/town/town-scm.test.ts @@ -0,0 +1,152 @@ +import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest'; +import { areThreadsBlocking, type SCMContext } from './town-scm'; +import { TownConfigSchema } from '../../types'; + +const THREADS = [ + { + isResolved: false, + comments: { nodes: [{ body: 'LGTM', author: { login: 'reviewer' } }] }, + }, +]; + +function makeCtx(config: Record) { + const aiRun = vi.fn(); + const ctx = { + env: { + AI: { run: aiRun }, + GASTOWN_AE: undefined, + KILO_API_URL: 'https://api.test', + }, + townId: 'town-1', + getTownConfig: async () => TownConfigSchema.parse({ town_id: 'town-1', ...config }), + } as unknown as SCMContext; + return { ctx, aiRun }; +} + +describe('areThreadsBlocking', () => { + beforeEach(() => { + vi.spyOn(console, 'warn').mockImplementation(() => {}); + vi.spyOn(console, 'log').mockImplementation(() => {}); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + vi.restoreAllMocks(); + }); + + it('uses the Kilo gateway when the configured model is direct BYOK', async () => { + const fetchMock = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ choices: [{ message: { content: '{"blocking": false}' } }] }), { + status: 200, + }) + ); + vi.stubGlobal('fetch', fetchMock); + + const { ctx, aiRun } = makeCtx({ + default_model: 'neuralwatt/glm-5.2-short', + kilocode_token: 'kilo-token', + }); + + expect(await areThreadsBlocking(ctx, THREADS)).toBe(false); + expect(aiRun).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(1); + + const [url, init] = fetchMock.mock.calls[0]; + expect(url).toBe('https://api.test/api/openrouter/chat/completions'); + expect(init.method).toBe('POST'); + expect(init.headers.Authorization).toBe('Bearer kilo-token'); + expect(init.headers['X-KiloCode-Feature']).toBe('gastown'); + expect(JSON.parse(init.body).model).toBe('neuralwatt/glm-5.2-short'); + }); + + it('prefers the refinery role model over the town default', async () => { + const fetchMock = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ choices: [{ message: { content: '{"blocking": false}' } }] }), { + status: 200, + }) + ); + vi.stubGlobal('fetch', fetchMock); + + const { ctx } = makeCtx({ + default_model: 'anthropic/claude-sonnet-4.6', + role_models: { refinery: 'zai-coding/glm-4.7' }, + kilocode_token: 'kilo-token', + }); + + expect(await areThreadsBlocking(ctx, THREADS)).toBe(false); + expect(JSON.parse(fetchMock.mock.calls[0][1].body).model).toBe('zai-coding/glm-4.7'); + }); + + it('sends the organization header for org towns', async () => { + const fetchMock = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ choices: [{ message: { content: '{"blocking": false}' } }] }), { + status: 200, + }) + ); + vi.stubGlobal('fetch', fetchMock); + + const { ctx } = makeCtx({ + default_model: 'neuralwatt/glm-5.2-short', + kilocode_token: 'kilo-token', + organization_id: 'org-1', + }); + + await areThreadsBlocking(ctx, THREADS); + expect(fetchMock.mock.calls[0][1].headers['X-KiloCode-OrganizationId']).toBe('org-1'); + }); + + it('falls back to Workers AI when the configured model is not direct BYOK', async () => { + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + + const { ctx, aiRun } = makeCtx({ + default_model: 'anthropic/claude-sonnet-4.6', + kilocode_token: 'kilo-token', + }); + aiRun.mockResolvedValue({ response: '{"blocking": false}' }); + + expect(await areThreadsBlocking(ctx, THREADS)).toBe(false); + expect(fetchMock).not.toHaveBeenCalled(); + expect(aiRun).toHaveBeenCalledWith('@cf/google/gemma-4-26b-a4b-it', expect.anything()); + }); + + it('falls back to Workers AI when no Kilo token is configured', async () => { + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + + const { ctx, aiRun } = makeCtx({ default_model: 'neuralwatt/glm-5.2-short' }); + aiRun.mockResolvedValue({ response: '{"blocking": false}' }); + + expect(await areThreadsBlocking(ctx, THREADS)).toBe(false); + expect(fetchMock).not.toHaveBeenCalled(); + expect(aiRun).toHaveBeenCalled(); + }); + + it('blocks without falling back to Workers AI when the gateway rejects the call', async () => { + const fetchMock = vi.fn().mockResolvedValue(new Response('payment required', { status: 402 })); + vi.stubGlobal('fetch', fetchMock); + + const { ctx, aiRun } = makeCtx({ + default_model: 'neuralwatt/glm-5.2-short', + kilocode_token: 'kilo-token', + }); + + expect(await areThreadsBlocking(ctx, THREADS)).toBe(true); + expect(aiRun).not.toHaveBeenCalled(); + }); + + it('blocks when the gateway response has no usable text', async () => { + const fetchMock = vi + .fn() + .mockResolvedValue(new Response(JSON.stringify({ choices: [] }), { status: 200 })); + vi.stubGlobal('fetch', fetchMock); + + const { ctx, aiRun } = makeCtx({ + default_model: 'neuralwatt/glm-5.2-short', + kilocode_token: 'kilo-token', + }); + + expect(await areThreadsBlocking(ctx, THREADS)).toBe(true); + expect(aiRun).not.toHaveBeenCalled(); + }); +}); diff --git a/services/gastown/src/dos/town/town-scm.ts b/services/gastown/src/dos/town/town-scm.ts index a077d713d3..6fd9c667fd 100644 --- a/services/gastown/src/dos/town/town-scm.ts +++ b/services/gastown/src/dos/town/town-scm.ts @@ -7,9 +7,13 @@ import { parseGitUrl, } from '../../util/platform-pr.util'; import { writeEvent } from '../../util/analytics.util'; +import { isDirectByokModelId } from '@kilocode/worker-utils/direct-byok-model'; const TOWN_LOG = '[town-scm]'; +/** Hardcoded Workers AI fallback model for the review-thread classifier. */ +const WORKERS_AI_CLASSIFIER_MODEL = '@cf/google/gemma-4-26b-a4b-it'; + export type SCMContext = { env: Env; townId: string; @@ -298,8 +302,52 @@ export async function checkPRStatus(ctx: SCMContext, prUrl: string): Promise { + const headers: Record = { + Authorization: `Bearer ${townConfig.kilocode_token}`, + 'Content-Type': 'application/json', + // Feature attribution for microdollar usage; 'gastown' is in FEATURE_VALUES + // (apps/web/src/lib/feature-detection.ts). + 'X-KiloCode-Feature': 'gastown', + }; + if (townConfig.organization_id) { + headers['X-KiloCode-OrganizationId'] = townConfig.organization_id; + } + + const response = await fetch(`${ctx.env.KILO_API_URL}/api/openrouter/chat/completions`, { + method: 'POST', + headers, + body: JSON.stringify({ + model, + messages: [{ role: 'user', content: prompt }], + max_tokens: 256, + temperature: 0, + }), + signal: AbortSignal.timeout(45_000), + }); + if (!response.ok) { + throw new Error( + `Kilo gateway returned ${response.status} for model ${model}: ${(await response.text()).slice(0, 300)}` + ); + } + return await response.json(); +} + +/** + * Determine if unresolved PR review threads contain blocking feedback that + * should prevent auto-merge. Direct-BYOK towns run the classifier on their + * configured model through the Kilo gateway; everyone else uses Workers AI. */ export async function areThreadsBlocking( ctx: SCMContext, @@ -331,20 +379,29 @@ Important: A comment is only NON-BLOCKING if it expresses approval or is purely Respond with ONLY a JSON object (no markdown, no explanation): { "blocking": true/false, "reason": "brief one-sentence explanation" }`; + const townConfig = await ctx.getTownConfig(); + // The refinery role owns the review flow, so its model is the one the user + // configured for reviews; fall back to the town default. + const configuredModel = townConfig.role_models?.refinery ?? townConfig.default_model; + const byokModel = + townConfig.kilocode_token && isDirectByokModelId(configuredModel) ? configuredModel : null; + const startTime = Date.now(); - const response: unknown = await ctx.env.AI.run('@cf/google/gemma-4-26b-a4b-it', { - messages: [{ role: 'user', content: prompt }], - max_tokens: 256, - temperature: 0, - chat_template_kwargs: { enable_thinking: false }, - }); + const response: unknown = byokModel + ? await classifyThreadsViaKiloGateway(ctx, townConfig, byokModel, prompt) + : await ctx.env.AI.run(WORKERS_AI_CLASSIFIER_MODEL, { + messages: [{ role: 'user', content: prompt }], + max_tokens: 256, + temperature: 0, + chat_template_kwargs: { enable_thinking: false }, + }); const durationMs = Date.now() - startTime; // Track the AI call via analytics event writeEvent(ctx.env, { event: 'api.external_request', townId: ctx.townId, - label: 'workers_ai_review_threads', + label: byokModel ? 'byok_review_threads' : 'workers_ai_review_threads', durationMs, }); From c6880dbbb2e8ab7e8d715eddfcfd5366a5749c14 Mon Sep 17 00:00:00 2001 From: maphew <486200+maphew@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:04:03 +0000 Subject: [PATCH 2/2] fix(gastown): disable reasoning for the BYOK thread classifier A reasoning-tagged BYOK model can spend the 256-token budget on a thinking trace and return no JSON content, which makes areThreadsBlocking() conservatively block auto-merge. Send reasoning: { enabled: false, effort: 'none' }, the gateway's "none" variant, and add tests for the refinery-override and gateway network-failure paths. Refs #4268 --- .../gastown/src/dos/town/town-scm.test.ts | 35 ++++++++++++++++++- services/gastown/src/dos/town/town-scm.ts | 12 +++++-- 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/services/gastown/src/dos/town/town-scm.test.ts b/services/gastown/src/dos/town/town-scm.test.ts index fe0e90dc94..271478f535 100644 --- a/services/gastown/src/dos/town/town-scm.test.ts +++ b/services/gastown/src/dos/town/town-scm.test.ts @@ -56,7 +56,11 @@ describe('areThreadsBlocking', () => { expect(init.method).toBe('POST'); expect(init.headers.Authorization).toBe('Bearer kilo-token'); expect(init.headers['X-KiloCode-Feature']).toBe('gastown'); - expect(JSON.parse(init.body).model).toBe('neuralwatt/glm-5.2-short'); + const body = JSON.parse(init.body); + expect(body.model).toBe('neuralwatt/glm-5.2-short'); + // Reasoning must be disabled or the model can spend the token budget on a + // thinking trace and return no JSON content. + expect(body.reasoning).toEqual({ enabled: false, effort: 'none' }); }); it('prefers the refinery role model over the town default', async () => { @@ -110,6 +114,22 @@ describe('areThreadsBlocking', () => { expect(aiRun).toHaveBeenCalledWith('@cf/google/gemma-4-26b-a4b-it', expect.anything()); }); + it('uses Workers AI when a managed refinery model overrides a BYOK default', async () => { + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + + const { ctx, aiRun } = makeCtx({ + default_model: 'neuralwatt/glm-5.2-short', + role_models: { refinery: 'anthropic/claude-sonnet-4.6' }, + kilocode_token: 'kilo-token', + }); + aiRun.mockResolvedValue({ response: '{"blocking": false}' }); + + expect(await areThreadsBlocking(ctx, THREADS)).toBe(false); + expect(fetchMock).not.toHaveBeenCalled(); + expect(aiRun).toHaveBeenCalled(); + }); + it('falls back to Workers AI when no Kilo token is configured', async () => { const fetchMock = vi.fn(); vi.stubGlobal('fetch', fetchMock); @@ -135,6 +155,19 @@ describe('areThreadsBlocking', () => { expect(aiRun).not.toHaveBeenCalled(); }); + it('blocks without falling back when the gateway request throws', async () => { + const fetchMock = vi.fn().mockRejectedValue(new Error('network down')); + vi.stubGlobal('fetch', fetchMock); + + const { ctx, aiRun } = makeCtx({ + default_model: 'neuralwatt/glm-5.2-short', + kilocode_token: 'kilo-token', + }); + + expect(await areThreadsBlocking(ctx, THREADS)).toBe(true); + expect(aiRun).not.toHaveBeenCalled(); + }); + it('blocks when the gateway response has no usable text', async () => { const fetchMock = vi .fn() diff --git a/services/gastown/src/dos/town/town-scm.ts b/services/gastown/src/dos/town/town-scm.ts index 6fd9c667fd..0b2d3955ca 100644 --- a/services/gastown/src/dos/town/town-scm.ts +++ b/services/gastown/src/dos/town/town-scm.ts @@ -305,8 +305,8 @@ export async function checkPRStatus(ctx: SCMContext, prUrl: string): Promise