Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
/**
* Tests for fixed Wiki Compose model id resolution.
*/
Comment thread
coderabbitai[bot] marked this conversation as resolved.
import { describe, expect, it, vi, beforeEach } from "vitest";
import {
resolveWikiComposeModelId,
WIKI_COMPOSE_MODEL_ID,
} from "../../../../agents/core/llm/wikiComposeModelId.js";

const mockDb = {
select: vi.fn(),
};

function chainLimit(rows: unknown[]) {
const chain = {
from: vi.fn().mockReturnThis(),
where: vi.fn().mockReturnThis(),
limit: vi.fn().mockResolvedValue(rows),
};
return chain;
}

beforeEach(() => {
vi.clearAllMocks();
});

describe("resolveWikiComposeModelId", () => {
it("returns the fixed id when the row is active and tier-accessible", async () => {
mockDb.select.mockReturnValueOnce(chainLimit([{ id: WIKI_COMPOSE_MODEL_ID }]));
const id = await resolveWikiComposeModelId("orchestrator", "free", mockDb as never);
expect(id).toBe(WIKI_COMPOSE_MODEL_ID);
});

it("returns the fixed id even when no DB row matches", async () => {
mockDb.select.mockReturnValueOnce(chainLimit([]));
const id = await resolveWikiComposeModelId("draft", "pro", mockDb as never);
expect(id).toBe("google:gemini-3.5-flash");
});

it("uses the same id for orchestrator and draft roles", async () => {
mockDb.select.mockReturnValue(chainLimit([{ id: WIKI_COMPOSE_MODEL_ID }]));
const orchestrator = await resolveWikiComposeModelId("orchestrator", "free", mockDb as never);
const draft = await resolveWikiComposeModelId("draft", "free", mockDb as never);
expect(orchestrator).toBe(draft);
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,11 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";

const { createZediChatModel } = vi.hoisted(() => ({ createZediChatModel: vi.fn() }));

vi.mock("../../../../../agents/core/llm/wikiComposeModelId.js", () => ({
WIKI_COMPOSE_MODEL_ID: "google:gemini-3.5-flash",
resolveWikiComposeModelId: vi.fn(async () => "google:gemini-3.5-flash"),
}));

vi.mock("../../../../../agents/core/llm/modelFactory.js", async () => {
const actual = await vi.importActual<
typeof import("../../../../../agents/core/llm/modelFactory.js")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,11 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";

const { createZediChatModel } = vi.hoisted(() => ({ createZediChatModel: vi.fn() }));

vi.mock("../../../../agents/core/llm/wikiComposeModelId.js", () => ({
WIKI_COMPOSE_MODEL_ID: "google:gemini-3.5-flash",
resolveWikiComposeModelId: vi.fn(async () => "google:gemini-3.5-flash"),
}));

vi.mock("../../../../agents/core/llm/modelFactory.js", async () => {
const actual = await vi.importActual<
typeof import("../../../../agents/core/llm/modelFactory.js")
Expand Down
11 changes: 5 additions & 6 deletions server/api/src/agents/core/composeBackendValidation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,12 @@
* Pre-flight BYOK checks for Wiki Compose session creation (#951).
* Wiki Compose セッション作成前の BYOK 事前チェック(#951)。
*
* Static env model ids are not validated here — provider matching is enforced at
* runtime via {@link resolveComposeModelId}. This function only verifies that
* the user has a stored credential when the graph will call an LLM.
* Wiki Compose uses a fixed Google model at runtime ({@link WIKI_COMPOSE_MODEL_ID}).
* BYOK sessions still need a stored credential when the graph calls an LLM; provider
* alignment with the fixed model is the caller's responsibility for now.
*
* 静的 env モデル id との provider 照合は行わない。実行時の
* `resolveComposeModelId` が provider 整合を担保する。本関数は LLM を呼ぶ
* グラフで credential が存在するかだけを確認する。
* Wiki Compose は実行時に Google 固定モデルを使う。BYOK 時は LLM 呼び出し前に
* credential の有無だけをここで確認する。

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Since Wiki Compose is now pinned to google:gemini-3.5-flash (via WIKI_COMPOSE_MODEL_ID), any BYOK backend other than google_byok (e.g., anthropic_byok or openai_byok) will pass this pre-flight check but fail at runtime with a BackendProviderMismatchError when createZediChatModel is called.

To prevent runtime crashes and provide a clear error message, we should update assertComposeBackendReady to reject non-Google BYOK backends when the graph uses WIKI_COMPOSE_MODEL_ID:

export async function assertComposeBackendReady(input: {
  backend: ExecutionBackend;
  graphId: string;
  userId: string;
  tier: UserTier;
  db: Database;
}): Promise<void> {
  if (!isUserByokBackend(input.backend)) return;

  const modelIds = getComposeModelIdsForGraph(input.graphId);
  if (modelIds.length === 0) return;

  const expectedProvider = backendToCredentialProvider(input.backend);

  // Reject non-Google BYOK backends early for Wiki Compose
  if (modelIds.includes(WIKI_COMPOSE_MODEL_ID) && expectedProvider !== "google") {
    throw new HTTPException(400, {
      message: `Wiki Compose is currently pinned to Google models. Backend "${input.backend}" is incompatible.`,
    });
  }

  const key = await getUserAiCredentialPlaintext(input.userId, expectedProvider, input.db);
  if (!key?.trim()) {
    throw new HTTPException(400, {
      message: `No API credential configured for backend "${input.backend}"`,
    });
  }
}

*/
import { HTTPException } from "hono/http-exception";
import type { Database, UserTier } from "../../types/index.js";
Expand Down
17 changes: 6 additions & 11 deletions server/api/src/agents/core/composeModelConfig.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,7 @@ import { WIKI_MAINTENANCE_GRAPH_ID } from "../graphs/wikiMaintenance/index.js";
import { getOrchestratorModelId } from "../subgraphs/research/nodes/planQueries.js";
import { RESEARCH_GRAPH_ID } from "../subgraphs/research/index.js";
import { INGEST_PLANNER_GRAPH_ID } from "../graphs/ingest/index.js";

const DRAFT_MODEL_ENV = "WIKI_COMPOSE_DRAFT_MODEL_ID";
const DRAFT_MODEL_FALLBACK = "claude-3-5-sonnet";

function getDraftModelId(): string {
return process.env[DRAFT_MODEL_ENV]?.trim() || DRAFT_MODEL_FALLBACK;
}
import { WIKI_COMPOSE_MODEL_ID } from "./llm/wikiComposeModelId.js";

/**
* Model row ids (`ai_models.id`) that a compose graph run will call via `createZediChatModel`.
Expand All @@ -23,11 +17,12 @@ export function getComposeModelIdsForGraph(graphId: string): string[] {
// Lint-only graph — no `createZediChatModel` calls; BYOK must not require orchestrator keys.
if (graphId === WIKI_MAINTENANCE_GRAPH_ID) return [];
if (graphId === WIKI_COMPOSE_GRAPH_ID) {
const orchestrator = getOrchestratorModelId();
const draft = getDraftModelId();
return orchestrator === draft ? [orchestrator] : [orchestrator, draft];
return [WIKI_COMPOSE_MODEL_ID];
}
if (graphId === RESEARCH_GRAPH_ID) {
return [WIKI_COMPOSE_MODEL_ID];
}
if (graphId === RESEARCH_GRAPH_ID || graphId === INGEST_PLANNER_GRAPH_ID) {
if (graphId === INGEST_PLANNER_GRAPH_ID) {
return [getOrchestratorModelId()];
}
return [getOrchestratorModelId()];
Expand Down
6 changes: 3 additions & 3 deletions server/api/src/agents/core/llm/resolveComposeModelId.ts
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
/**
* Resolve `ai_models.id` for Wiki Compose nodes so BYOK backends use a matching provider.
* Resolve `ai_models.id` for non–Wiki Compose graphs (e.g. ingest planner).
* Wiki Compose uses {@link resolveWikiComposeModelId} instead.
*
* Wiki Compose の LLM ノード用 model id 解決。BYOK backend では provider が一致する
* active モデルを選び、`BackendProviderMismatchError` でセッション全体が落ちるのを防ぐ。
* Wiki Compose 以外(ingest planner 等)向けの model id 解決。
*/
import { and, asc, eq } from "drizzle-orm";
import { aiModels } from "../../../schema/index.js";
Expand Down
54 changes: 54 additions & 0 deletions server/api/src/agents/core/llm/wikiComposeModelId.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
/**
* Fixed Wiki Compose model id (temporary until per-role / per-backend selection returns).
* Wiki Compose 用の固定モデル id(将来ロール別選択に戻すまでの暫定)。
*/
import { and, eq } from "drizzle-orm";
import { aiModels } from "../../../schema/index.js";
import type { Database, UserTier } from "../../../types/index.js";
import type { ComposeModelRole } from "./resolveComposeModelId.js";

/**
* `ai_models.id` used by every Wiki Compose LLM node and the `web_search` tool.
* すべての Wiki Compose LLM ノードと `web_search` ツールが使う `ai_models.id`。
*/
export const WIKI_COMPOSE_MODEL_ID = "google:gemini-3.5-flash" as const;

function tierFilter(tier: UserTier) {
if (tier === "pro") return undefined;
return eq(aiModels.tierRequired, "free");
}

/**
* When the fixed row exists and is accessible, return it; otherwise still return
* {@link WIKI_COMPOSE_MODEL_ID} so callers fail consistently in `validateModelAccess`.
* 行が active かつ tier 的に使えるときだけ DB 上の id を返し、それ以外は固定 id を返す。
*/
async function fixedModelIdIfAccessible(db: Database, tier: UserTier): Promise<string> {
const tierClause = tierFilter(tier);
const [row] = await db
.select({ id: aiModels.id })
.from(aiModels)
.where(
and(
eq(aiModels.id, WIKI_COMPOSE_MODEL_ID),
eq(aiModels.isActive, true),
...(tierClause ? [tierClause] : []),
),
)
.limit(1);
return row?.id ?? WIKI_COMPOSE_MODEL_ID;
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The database query to resolve the active and accessible fixed model ID is duplicated in resolveWebSearchModel.ts. We can extract this query into a shared helper function resolveActiveWikiComposeModelId to improve maintainability and reuse it across both files.

export async function resolveActiveWikiComposeModelId(
  db: Database,
  tier: UserTier,
): Promise<string | null> {
  const tierClause = tierFilter(tier);
  const [row] = await db
    .select({ id: aiModels.id })
    .from(aiModels)
    .where(
      and(
        eq(aiModels.id, WIKI_COMPOSE_MODEL_ID),
        eq(aiModels.isActive, true),
        ...(tierClause ? [tierClause] : []),
      ),
    )
    .limit(1);
  return row?.id ?? null;
}

async function fixedModelIdIfAccessible(db: Database, tier: UserTier): Promise<string> {
  return (await resolveActiveWikiComposeModelId(db, tier)) ?? WIKI_COMPOSE_MODEL_ID;
}


/**
* Resolve the model row id for Wiki Compose orchestrator / draft / research nodes.
* `role` is accepted for API stability; all roles map to {@link WIKI_COMPOSE_MODEL_ID} for now.
*
* Wiki Compose の model id 解決。現時点では role に関わらず gemini-3.5-flash 固定。
*/
export async function resolveWikiComposeModelId(
_role: ComposeModelRole,
_tier: UserTier,
db: Database,
): Promise<string> {
return fixedModelIdIfAccessible(db, _tier);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Reject incompatible BYOK backends before pinning model

resolveWikiComposeModelId now always resolves to google:gemini-3.5-flash regardless of the session backend, so user_openai / user_anthropic sessions can still be created but will fail at runtime when createZediChatModel enforces provider/backend matching (BackendProviderMismatchError). This is a regression from the previous backend-aware resolution path and will break valid BYOK users unless backend compatibility is checked up front or model resolution is backend-aware again.

Useful? React with 👍 / 👎.

}
25 changes: 20 additions & 5 deletions server/api/src/agents/core/tools/resolveWebSearchModel.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,14 @@
* `useGoogleSearch` for Google) を呼ぶため、Anthropic-only な選択では成立しない。
* 本ヘルパは次の優先順で model を選ぶ:
*
* 1. `process.env.WIKI_COMPOSE_WEB_SEARCH_MODEL_ID` (explicit override; `ai_models.id`)
* 1. {@link WIKI_COMPOSE_MODEL_ID} when active and tier-accessible
* 2. `process.env.WIKI_COMPOSE_WEB_SEARCH_MODEL_ID` (explicit override; `ai_models.id`)
* — 必ず active かつ tier 通過することを DB 側で確認する(coderabbit review #956:
* 不正な override で `createZediChatModel` が失敗してエラー envelope になる
* のを防ぐ)。
* 2. `ai_models` の active な OpenAI モデルで最安 (`input_cost_units` ASC, `output_cost_units` ASC)
* 3. `ai_models` の active な Google モデルで最安
* 4. 何も無ければ `null` を返す(ツール側は empty result + note を返す)。
* 3. `ai_models` の active な OpenAI モデルで最安 (`input_cost_units` ASC, `output_cost_units` ASC)
* 4. `ai_models` の active な Google モデルで最安
* 5. 何も無ければ `null` を返す(ツール側は empty result + note を返す)。
*
* Returns the `ai_models.id` so `createZediChatModel({ modelId })` can validate
* tier access and resolve the API key uniformly. Centralising the choice in one
Expand All @@ -25,6 +26,7 @@
import { and, asc, eq, inArray } from "drizzle-orm";
import { aiModels } from "../../../schema/index.js";
import type { Database, UserTier } from "../../../types/index.js";
import { WIKI_COMPOSE_MODEL_ID } from "../llm/wikiComposeModelId.js";

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Import the newly extracted resolveActiveWikiComposeModelId helper from wikiComposeModelId.ts.

Suggested change
import { WIKI_COMPOSE_MODEL_ID } from "../llm/wikiComposeModelId.js";
import { WIKI_COMPOSE_MODEL_ID, resolveActiveWikiComposeModelId } from "../llm/wikiComposeModelId.js";


const ENV_OVERRIDE = "WIKI_COMPOSE_WEB_SEARCH_MODEL_ID";

Expand All @@ -48,6 +50,20 @@ export async function resolveWebSearchModelId(
db: Database,
tier: UserTier,
): Promise<string | null> {
const tierClause = tierFilter(tier);
const [fixedRow] = await db
.select({ id: aiModels.id })
.from(aiModels)
.where(
and(
eq(aiModels.id, WIKI_COMPOSE_MODEL_ID),
eq(aiModels.isActive, true),
...(tierClause ? [tierClause] : []),
),
)
.limit(1);
if (fixedRow) return fixedRow.id;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Reuse the extracted resolveActiveWikiComposeModelId helper to avoid duplicating the database query logic.

  const fixedId = await resolveActiveWikiComposeModelId(db, tier);
  if (fixedId) return fixedId;

  const tierClause = tierFilter(tier);


const override = process.env[ENV_OVERRIDE]?.trim();
if (override) {
// Validate the override before returning: it must be active and
Expand All @@ -73,7 +89,6 @@ export async function resolveWebSearchModelId(
// override が使えない場合は通常検索にフォールバックする。
}

const tierClause = tierFilter(tier);
const rows = await db
.select({
id: aiModels.id,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ import { randomUUID } from "node:crypto";
import { z } from "zod";
import { composeContentLocaleInstruction } from "../../../core/composeLocale.js";
import { createZediChatModel } from "../../../core/llm/modelFactory.js";
import { resolveComposeModelId } from "../../../core/llm/resolveComposeModelId.js";
import { resolveWikiComposeModelId } from "../../../core/llm/wikiComposeModelId.js";
import { getGraphContext } from "../../../subgraphs/research/nodes/shared/getGraphContext.js";
import { loadPageSnapshot } from "./shared/loadPageSnapshot.js";
import { dispatchComposePhase } from "./shared/dispatch.js";
Expand Down Expand Up @@ -117,7 +117,7 @@ export async function briefDialogue(
// セッション開始時に 1 度だけ読み、以後は state を参照する。
const snapshot = state.pageSnapshot ?? (await loadPageSnapshot(ctx.db, ctx.pageId));

const modelId = await resolveComposeModelId("orchestrator", ctx.backend, ctx.tier, ctx.db);
const modelId = await resolveWikiComposeModelId("orchestrator", ctx.tier, ctx.db);
const model = await createZediChatModel({
modelId,
userId: ctx.userId,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
import type { LangGraphRunnableConfig } from "@langchain/langgraph";
import { composeContentLocaleInstruction } from "../../../core/composeLocale.js";
import { createZediChatModel } from "../../../core/llm/modelFactory.js";
import { resolveComposeModelId } from "../../../core/llm/resolveComposeModelId.js";
import { resolveWikiComposeModelId } from "../../../core/llm/wikiComposeModelId.js";
import { getGraphContext } from "../../../subgraphs/research/nodes/shared/getGraphContext.js";
import { dispatchComposePhase, dispatchComposeSection } from "./shared/dispatch.js";
import type { WikiComposeStateType, WikiComposeStateUpdate } from "../state.js";
Expand Down Expand Up @@ -127,7 +127,7 @@ export async function draftSections(
return { draftedSections: [], phase: "draft:completed" };
}

const modelId = await resolveComposeModelId("draft", ctx.backend, ctx.tier, ctx.db);
const modelId = await resolveWikiComposeModelId("draft", ctx.tier, ctx.db);
const model = await createZediChatModel({
modelId,
userId: ctx.userId,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import {
structureDialogueFallbackOutline,
} from "../../../core/composeLocale.js";
import { createZediChatModel } from "../../../core/llm/modelFactory.js";
import { resolveComposeModelId } from "../../../core/llm/resolveComposeModelId.js";
import { resolveWikiComposeModelId } from "../../../core/llm/wikiComposeModelId.js";
import { getGraphContext } from "../../../subgraphs/research/nodes/shared/getGraphContext.js";
import { dispatchComposePhase } from "./shared/dispatch.js";
import type { WikiComposeStateType, WikiComposeStateUpdate } from "../state.js";
Expand Down Expand Up @@ -83,7 +83,7 @@ export async function structureDialogue(

await dispatchComposePhase({ phase: "structure", status: "entered" }, config);

const modelId = await resolveComposeModelId("orchestrator", ctx.backend, ctx.tier, ctx.db);
const modelId = await resolveWikiComposeModelId("orchestrator", ctx.tier, ctx.db);
const model = await createZediChatModel({
modelId,
userId: ctx.userId,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ import type { LangGraphRunnableConfig } from "@langchain/langgraph";
import { z } from "zod";
import { composeContentLocaleInstruction } from "../../../core/composeLocale.js";
import { createZediChatModel } from "../../../core/llm/modelFactory.js";
import { resolveComposeModelId } from "../../../core/llm/resolveComposeModelId.js";
import { resolveWikiComposeModelId } from "../../../core/llm/wikiComposeModelId.js";
import { getGraphContext } from "./shared/getGraphContext.js";
import { dispatchResearchEvaluation } from "./shared/dispatchSseCustom.js";
import type { ResearchLoopStateType, ResearchLoopStateUpdate } from "../state.js";
Expand Down Expand Up @@ -72,7 +72,7 @@ export async function evaluateSufficiency(
): Promise<ResearchLoopStateUpdate> {
const ctx = getGraphContext(config);

const modelId = await resolveComposeModelId("orchestrator", ctx.backend, ctx.tier, ctx.db);
const modelId = await resolveWikiComposeModelId("orchestrator", ctx.tier, ctx.db);
const model = await createZediChatModel({
modelId,
userId: ctx.userId,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,15 @@ import { randomUUID } from "node:crypto";
import { z } from "zod";
import { composeContentLocaleInstruction } from "../../../core/composeLocale.js";
import { createZediChatModel } from "../../../core/llm/modelFactory.js";
import { resolveComposeModelId } from "../../../core/llm/resolveComposeModelId.js";
import { resolveWikiComposeModelId } from "../../../core/llm/wikiComposeModelId.js";
import { getGraphContext } from "./shared/getGraphContext.js";
import { dispatchResearchIteration } from "./shared/dispatchSseCustom.js";
import type { ResearchLoopStateType, ResearchLoopStateUpdate } from "../state.js";
import type { PlannedQuery, Source } from "../types.js";

/**
* @deprecated Use {@link resolveComposeModelId} with graph context. Kept for tests importing the symbol.
* @deprecated Use {@link resolveWikiComposeModelId} for Wiki Compose graphs.
* Kept for ingest planner BYOK preflight ({@link getComposeModelIdsForGraph}).
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
*/
export function getOrchestratorModelId(): string {
return process.env.WIKI_COMPOSE_ORCHESTRATOR_MODEL_ID?.trim() || "claude-3-5-haiku";
Expand Down Expand Up @@ -100,7 +101,7 @@ export async function planQueries(
// maxIterations は既存 state を優先しつつ 1..5 にクランプ。
const maxIterations = clampMaxIterations(state.maxIterations ?? 3);

const modelId = await resolveComposeModelId("orchestrator", ctx.backend, ctx.tier, ctx.db);
const modelId = await resolveWikiComposeModelId("orchestrator", ctx.tier, ctx.db);
const model = await createZediChatModel({
modelId,
userId: ctx.userId,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import type { LangGraphRunnableConfig } from "@langchain/langgraph";
import { randomUUID } from "node:crypto";
import { composeContentLocaleInstruction } from "../../../core/composeLocale.js";
import { createZediChatModel } from "../../../core/llm/modelFactory.js";
import { resolveComposeModelId } from "../../../core/llm/resolveComposeModelId.js";
import { resolveWikiComposeModelId } from "../../../core/llm/wikiComposeModelId.js";
import { getGraphContext } from "./shared/getGraphContext.js";
import { dispatchResearchIteration } from "./shared/dispatchSseCustom.js";
import { planQueriesSchema } from "./planQueries.js";
Expand Down Expand Up @@ -62,7 +62,7 @@ export async function refineQueries(
): Promise<ResearchLoopStateUpdate> {
const ctx = getGraphContext(config);

const modelId = await resolveComposeModelId("orchestrator", ctx.backend, ctx.tier, ctx.db);
const modelId = await resolveWikiComposeModelId("orchestrator", ctx.tier, ctx.db);
const model = await createZediChatModel({
modelId,
userId: ctx.userId,
Expand Down
Loading