From 04844c0b64fa96b96c23026d17f9e1260e5ec6e1 Mon Sep 17 00:00:00 2001 From: tomsmith8 Date: Wed, 2 Sep 2026 14:24:57 +0000 Subject: [PATCH 1/2] Generated with Hive: Add xAI as first-class model provider with direct XAI_API_KEY authentication --- env.example | 6 + .../migration.sql | 20 ++ prisma/schema.prisma | 1 + .../integration/api/admin/llm-models.test.ts | 199 ++++++++++++++++++ .../integration/api/llm-models.test.ts | 71 +++++++ .../unit/api/user/preferences.test.ts | 54 +++++ src/__tests__/unit/lib/ai/models.test.ts | 62 +++++- .../unit/lib/ai/runCanvasAgent-xai.test.ts | 117 ++++++++++ src/app/admin/llm-models/LlmModelsTable.tsx | 3 +- src/app/api/admin/llm-models/[id]/route.ts | 130 ++++++++++-- src/app/api/admin/llm-models/route.ts | 129 +++++++++--- src/app/api/agent/route.ts | 83 ++++++-- src/app/api/llm-models/route.ts | 28 ++- src/app/api/user/preferences/route.ts | 32 ++- .../[slug]/legal/benchmarks/run/route.ts | 12 ++ .../CanvasAgentSettingsPopover.tsx | 13 +- src/components/legal/TaskDetailsModal.tsx | 7 +- src/config/env.ts | 8 + src/lib/ai/models.ts | 5 + src/lib/ai/runCanvasAgent.ts | 17 +- src/services/task-workflow.ts | 116 ++++++++-- src/utils/mockSetup.ts | 14 ++ 22 files changed, 1044 insertions(+), 83 deletions(-) create mode 100644 prisma/migrations/20260902134634_add_xai_provider/migration.sql create mode 100644 src/__tests__/unit/lib/ai/runCanvasAgent-xai.test.ts diff --git a/env.example b/env.example index 9270ebbaf5..9e1f3c01d6 100644 --- a/env.example +++ b/env.example @@ -149,6 +149,12 @@ FEATURE_WHITEBOARD_STAKWORK_POSITIONING=false ANTHROPIC_API_KEY="your-anthropic-api-key-here" # OpenAI API Key (optional, for future features) # OPENAI_API_KEY="" +# xAI API Key — direct credential for Grok models (provider "XAI" in the +# LlmModel admin table). Without this, xAI/Grok rows are filtered out of +# every model picker (see /api/llm-models) and dispatch falls back to a +# "model provider not configured" error rather than routing through +# OpenRouter. Get your key from: https://console.x.ai/ +# XAI_API_KEY="your-xai-api-key-here" # Payment Processing (for bounty system) # STRIPE_SECRET_KEY="" diff --git a/prisma/migrations/20260902134634_add_xai_provider/migration.sql b/prisma/migrations/20260902134634_add_xai_provider/migration.sql new file mode 100644 index 0000000000..b75c6ff6c8 --- /dev/null +++ b/prisma/migrations/20260902134634_add_xai_provider/migration.sql @@ -0,0 +1,20 @@ +-- AlterEnum +-- Add XAI as a first-class LlmProvider value (idempotent). +-- Only alter if the enum type exists (handles shadow DB scenario) and only +-- add the label if it isn't already present (handles preview-branch re-runs +-- and avoids Postgres 42710 "enum label already exists"). +-- Mirrors the pattern in 20260124103308_add_pod_status_and_soft_delete. +-- +-- One-way: Postgres cannot DROP an enum value. Walking this feature back +-- just means leaving an orphaned, unused label — no down-migration needed. +DO $$ +DECLARE enum_oid OID; +BEGIN + SELECT oid INTO enum_oid FROM pg_type WHERE typname = 'LlmProvider'; + + IF enum_oid IS NOT NULL AND NOT EXISTS ( + SELECT 1 FROM pg_enum WHERE enumlabel = 'XAI' AND enumtypid = enum_oid + ) THEN + ALTER TYPE "LlmProvider" ADD VALUE 'XAI'; + END IF; +END $$; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index c91796a840..5f4228ba63 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -1703,6 +1703,7 @@ enum LlmProvider { ANTHROPIC OPENAI AWS_BEDROCK + XAI OTHER } diff --git a/src/__tests__/integration/api/admin/llm-models.test.ts b/src/__tests__/integration/api/admin/llm-models.test.ts index 3b41688788..a9639aa594 100644 --- a/src/__tests__/integration/api/admin/llm-models.test.ts +++ b/src/__tests__/integration/api/admin/llm-models.test.ts @@ -462,5 +462,204 @@ describe("Admin LLM Models API", () => { expect(response.status).toBe(401); }); + + it("should return 403 when token-authenticated PATCH tries to set isPublic", async () => { + const model = await createTestLlmModel({ name: "sync-patch-isPublic-blocked", provider: "OPENAI" }); + const request = createPatchRequestWithApiToken( + `/api/admin/llm-models/${model.id}`, + TEST_API_TOKEN, + { isPublic: true }, + ); + const { PATCH } = await import("@/app/api/admin/llm-models/[id]/route"); + const response = await PATCH(request, { + params: Promise.resolve({ id: model.id }), + }); + + expect(response.status).toBe(403); + const dbRecord = await db.llmModel.findUnique({ where: { id: model.id } }); + expect(dbRecord?.isPublic).toBe(false); + }); + + it("should return 403 when token-authenticated PATCH tries to set isPlanDefault or isTaskDefault", async () => { + const model = await createTestLlmModel({ name: "sync-patch-defaults-blocked", provider: "OPENAI" }); + const request = createPatchRequestWithApiToken( + `/api/admin/llm-models/${model.id}`, + TEST_API_TOKEN, + { isPlanDefault: true, isTaskDefault: true }, + ); + const { PATCH } = await import("@/app/api/admin/llm-models/[id]/route"); + const response = await PATCH(request, { + params: Promise.resolve({ id: model.id }), + }); + + expect(response.status).toBe(403); + }); + + it("should allow a SUPER_ADMIN session to set isPublic/isPlanDefault/isTaskDefault", async () => { + const model = await createTestLlmModel({ name: "session-patch-defaults-allowed", provider: "OPENAI" }); + const request = createAuthenticatedPatchRequest( + `/api/admin/llm-models/${model.id}`, + { isPublic: true, isPlanDefault: true }, + superAdminUser, + ); + const { PATCH } = await import("@/app/api/admin/llm-models/[id]/route"); + const response = await PATCH(request, { + params: Promise.resolve({ id: model.id }), + }); + + expect(response.status).toBe(200); + const data = await response.json(); + expect(data.model.isPublic).toBe(true); + expect(data.model.isPlanDefault).toBe(true); + }); + }); + + describe("name/providerLabel validation", () => { + it("should reject a name containing a slash on POST", async () => { + const request = createAuthenticatedPostRequest( + "/api/admin/llm-models", + superAdminUser, + { name: "anthropic/claude-x", provider: "XAI", inputPricePer1M: 1, outputPricePer1M: 2 }, + ); + const { POST } = await import("@/app/api/admin/llm-models/route"); + const response = await POST(request); + + expect(response.status).toBe(400); + }); + + it("should reject a name containing a slash on PATCH", async () => { + const model = await createTestLlmModel({ name: "safe-name-model", provider: "OPENAI" }); + const request = createAuthenticatedPatchRequest( + `/api/admin/llm-models/${model.id}`, + { name: "openai/gpt-x" }, + superAdminUser, + ); + const { PATCH } = await import("@/app/api/admin/llm-models/[id]/route"); + const response = await PATCH(request, { + params: Promise.resolve({ id: model.id }), + }); + + expect(response.status).toBe(400); + }); + + it("should reject a providerLabel containing a slash", async () => { + const request = createAuthenticatedPostRequest( + "/api/admin/llm-models", + superAdminUser, + { + name: "custom-model", + provider: "OTHER", + providerLabel: "Foo/Bar", + inputPricePer1M: 1, + outputPricePer1M: 2, + }, + ); + const { POST } = await import("@/app/api/admin/llm-models/route"); + const response = await POST(request); + + expect(response.status).toBe(400); + }); + + it("should accept a valid XAI model name", async () => { + const request = createAuthenticatedPostRequest( + "/api/admin/llm-models", + superAdminUser, + { name: "grok-4", provider: "XAI", inputPricePer1M: 3, outputPricePer1M: 15 }, + ); + const { POST } = await import("@/app/api/admin/llm-models/route"); + const response = await POST(request); + + expect(response.status).toBe(201); + const data = await response.json(); + expect(data.model.provider).toBe("XAI"); + }); + }); + + describe("Duplicate name conflict", () => { + it("should return 409 with the existing row's id when creating a duplicate name", async () => { + const existing = await createTestLlmModel({ name: "dup-model", provider: "OPENAI" }); + + const request = createAuthenticatedPostRequest( + "/api/admin/llm-models", + superAdminUser, + { name: "dup-model", provider: "ANTHROPIC", inputPricePer1M: 1, outputPricePer1M: 2 }, + ); + const { POST } = await import("@/app/api/admin/llm-models/route"); + const response = await POST(request); + + expect(response.status).toBe(409); + const data = await response.json(); + expect(data.existingId).toBe(existing.id); + }); + }); + + describe("Default-flip atomicity", () => { + it("should leave exactly one isPlanDefault: true row after sequential flips", async () => { + const first = await createTestLlmModel({ name: "default-flip-a", provider: "OPENAI", isPlanDefault: true } as never); + + const request = createAuthenticatedPostRequest( + "/api/admin/llm-models", + superAdminUser, + { + name: "default-flip-b", + provider: "ANTHROPIC", + inputPricePer1M: 1, + outputPricePer1M: 2, + isPlanDefault: true, + }, + ); + const { POST } = await import("@/app/api/admin/llm-models/route"); + await POST(request); + + const defaults = await db.llmModel.findMany({ + where: { isPlanDefault: true, id: { in: [first.id] } }, + }); + // The original row should have had its default cleared by the + // transactional flip triggered by the new row's create. + expect(defaults).toHaveLength(0); + + const allDefaults = await db.llmModel.findMany({ where: { isPlanDefault: true } }); + const relevant = allDefaults.filter((m) => ["default-flip-a", "default-flip-b"].includes(m.name)); + expect(relevant).toHaveLength(1); + expect(relevant[0].name).toBe("default-flip-b"); + }); + }); + + describe("Sync-revert regression (provider/providerLabel excluded from batch-upsert update)", () => { + beforeEach(() => { + process.env.API_TOKEN = TEST_API_TOKEN; + }); + + it("should leave provider unchanged on batch upsert of an existing XAI row, while still updating pricing", async () => { + await createTestLlmModel({ + name: "grok-sync-model", + provider: "XAI", + inputPricePer1M: 3.0, + outputPricePer1M: 15.0, + }); + + // Simulate the nightly sync workflow re-classifying the row to OTHER — + // this must NOT stick; only pricing should update. + const request = createRequestWithApiToken("/api/admin/llm-models", TEST_API_TOKEN, { + models: [ + { + name: "grok-sync-model", + provider: "OTHER", + providerLabel: "OpenRouter", + inputPricePer1M: 9.99, + outputPricePer1M: 29.99, + }, + ], + }); + const { POST } = await import("@/app/api/admin/llm-models/route"); + const response = await POST(request); + expect(response.status).toBe(201); + + const dbRecord = await db.llmModel.findUnique({ where: { name: "grok-sync-model" } }); + expect(dbRecord?.provider).toBe("XAI"); + expect(dbRecord?.providerLabel).toBeNull(); + expect(dbRecord?.inputPricePer1M).toBe(9.99); + expect(dbRecord?.outputPricePer1M).toBe(29.99); + }); }); }); diff --git a/src/__tests__/integration/api/llm-models.test.ts b/src/__tests__/integration/api/llm-models.test.ts index 0155e1dca6..632e57178f 100644 --- a/src/__tests__/integration/api/llm-models.test.ts +++ b/src/__tests__/integration/api/llm-models.test.ts @@ -32,9 +32,22 @@ function createGetRequest(token?: string) { describe("GET /api/llm-models - Integration Tests", () => { let seededModels: LlmModel[] = []; let testUser: User; + const originalXaiKey = process.env.XAI_API_KEY; + const originalAnthropicKey = process.env.ANTHROPIC_API_KEY; + const originalOpenaiKey = process.env.OPENAI_API_KEY; + const originalGoogleKey = process.env.GOOGLE_API_KEY; beforeEach(async () => { process.env.API_TOKEN = VALID_API_TOKEN; + // The provider-key-availability filter (isProviderKeyConfigured) means + // every pre-existing test that seeds OPENAI/ANTHROPIC/GOOGLE models + // needs those providers "configured" to see them in the response — + // mirroring production, where these are always set. Only the new + // "Provider key availability filter" tests below manipulate these + // per-case. + process.env.OPENAI_API_KEY = "test-openai-key"; + process.env.ANTHROPIC_API_KEY = "test-anthropic-key"; + process.env.GOOGLE_API_KEY = "test-google-key"; testUser = await createTestUser(); @@ -79,6 +92,10 @@ describe("GET /api/llm-models - Integration Tests", () => { }); await db.user.deleteMany({ where: { id: testUser.id } }); seededModels = []; + if (originalXaiKey === undefined) delete process.env.XAI_API_KEY; + else process.env.XAI_API_KEY = originalXaiKey; + if (originalAnthropicKey === undefined) delete process.env.ANTHROPIC_API_KEY; + else process.env.ANTHROPIC_API_KEY = originalAnthropicKey; }); describe("Authentication", () => { @@ -283,4 +300,58 @@ describe("GET /api/llm-models - Integration Tests", () => { expect(claude.providerLabel).toBe("Anthropic"); }); }); + + describe("Provider key availability filter", () => { + test("excludes an XAI model when XAI_API_KEY is not set", async () => { + delete process.env.XAI_API_KEY; + const xaiModel = await createTestLlmModel({ + name: "grok-4", + provider: "XAI", + inputPricePer1M: 3.0, + outputPricePer1M: 15.0, + dateEnd: null, + isPublic: true, + }); + seededModels.push(xaiModel); + + const request = createGetRequest(VALID_API_TOKEN); + const response = await GET(request as any); + const data = await response.json(); + const returnedIds = data.models.map((m: { id: string }) => m.id); + + expect(returnedIds).not.toContain(xaiModel.id); + }); + + test("includes an XAI model when XAI_API_KEY is set", async () => { + process.env.XAI_API_KEY = "test-xai-key"; + const xaiModel = await createTestLlmModel({ + name: "grok-4-set", + provider: "XAI", + inputPricePer1M: 3.0, + outputPricePer1M: 15.0, + dateEnd: null, + isPublic: true, + }); + seededModels.push(xaiModel); + + const request = createGetRequest(VALID_API_TOKEN); + const response = await GET(request as any); + const data = await response.json(); + const returnedIds = data.models.map((m: { id: string }) => m.id); + + expect(returnedIds).toContain(xaiModel.id); + }); + + test("does not affect an existing provider's models when its key is set (regression)", async () => { + process.env.ANTHROPIC_API_KEY = "test-anthropic-key"; + + const request = createGetRequest(VALID_API_TOKEN); + const response = await GET(request as any); + const data = await response.json(); + const returnedIds = data.models.map((m: { id: string }) => m.id); + const claudeModel = seededModels.find((m) => m.name === "claude-3-5-sonnet")!; + + expect(returnedIds).toContain(claudeModel.id); + }); + }); }); diff --git a/src/__tests__/unit/api/user/preferences.test.ts b/src/__tests__/unit/api/user/preferences.test.ts index a197e90bc4..1c98306b56 100644 --- a/src/__tests__/unit/api/user/preferences.test.ts +++ b/src/__tests__/unit/api/user/preferences.test.ts @@ -13,6 +13,7 @@ vi.mock("@/lib/auth/nextauth", () => ({ const mockUserFindUnique = vi.fn(); const mockUserUpdate = vi.fn(); +const mockLlmModelFindMany = vi.fn(); vi.mock("@/lib/db", () => ({ db: { @@ -20,6 +21,9 @@ vi.mock("@/lib/db", () => ({ findUnique: (...args: unknown[]) => mockUserFindUnique(...args), update: (...args: unknown[]) => mockUserUpdate(...args), }, + llmModel: { + findMany: (...args: unknown[]) => mockLlmModelFindMany(...args), + }, }, })); @@ -110,6 +114,56 @@ describe("GET /api/user/preferences", () => { const res = await GET(); expect(res.status).toBe(401); }); + + test("falls back to null chatAgentModel when the stored preference no longer matches any catalog row", async () => { + mockUserFindUnique.mockResolvedValue({ + canvasAutonomousTurns: false, + chatAgentModel: "openrouter/grok-4", + timezone: "UTC", + dailyRecapEnabled: true, + }); + mockLlmModelFindMany.mockResolvedValue([]); + + const res = await GET(); + const body = await res.json(); + + expect(res.status).toBe(200); + expect(body.chatAgentModel).toBeNull(); + }); + + test("keeps chatAgentModel when it still matches a public, unexpired catalog row", async () => { + mockUserFindUnique.mockResolvedValue({ + canvasAutonomousTurns: false, + chatAgentModel: "xai/grok-4", + timezone: "UTC", + dailyRecapEnabled: true, + }); + mockLlmModelFindMany.mockResolvedValue([ + { id: "m1", name: "grok-4", provider: "XAI", providerLabel: null, isPlanDefault: false, isTaskDefault: false }, + ]); + + const res = await GET(); + const body = await res.json(); + + expect(res.status).toBe(200); + expect(body.chatAgentModel).toBe("xai/grok-4"); + }); + + test("passes null through unchanged (no catalog lookup needed)", async () => { + mockUserFindUnique.mockResolvedValue({ + canvasAutonomousTurns: false, + chatAgentModel: null, + timezone: "UTC", + dailyRecapEnabled: true, + }); + + const res = await GET(); + const body = await res.json(); + + expect(res.status).toBe(200); + expect(body.chatAgentModel).toBeNull(); + expect(mockLlmModelFindMany).not.toHaveBeenCalled(); + }); }); describe("PATCH /api/user/preferences — timezone", () => { diff --git a/src/__tests__/unit/lib/ai/models.test.ts b/src/__tests__/unit/lib/ai/models.test.ts index 6fd86f57af..b4d942287e 100644 --- a/src/__tests__/unit/lib/ai/models.test.ts +++ b/src/__tests__/unit/lib/ai/models.test.ts @@ -1,6 +1,6 @@ // @vitest-environment jsdom import { describe, test, expect, vi, beforeEach, afterEach } from "vitest"; -import { isValidModel, getApiKeyForModel, VALID_MODELS, PROVIDER_API_KEY_ENV_VARS, getStoredPlanModelPreference, setStoredPlanModelPreference, PLAN_MODEL_PREFERENCE_KEY } from "@/lib/ai/models"; +import { isValidModel, getApiKeyForModel, getModelValue, VALID_MODELS, PROVIDER_API_KEY_ENV_VARS, PROVIDER_DISPLAY_LABELS, getStoredPlanModelPreference, setStoredPlanModelPreference, PLAN_MODEL_PREFERENCE_KEY } from "@/lib/ai/models"; describe("models", () => { describe("isValidModel", () => { @@ -29,6 +29,10 @@ describe("models", () => { expect(isValidModel("unknown/some-model")).toBe(false); }); + test("returns true for xai/name format", () => { + expect(isValidModel("xai/grok-4")).toBe(true); + }); + test("returns false for non-string values", () => { expect(isValidModel(null)).toBe(false); expect(isValidModel(undefined)).toBe(false); @@ -99,6 +103,62 @@ describe("models", () => { process.env.GOOGLE_API_KEY = "test-google-key"; expect(getApiKeyForModel("gemini")).toBe("test-google-key"); }); + + test("returns XAI_API_KEY for provider/name format (xai/...)", () => { + process.env.XAI_API_KEY = "test-xai-key"; + expect(getApiKeyForModel("xai/grok-4")).toBe("test-xai-key"); + }); + + test("returns undefined for xai/... when XAI_API_KEY is not set", () => { + delete process.env.XAI_API_KEY; + expect(getApiKeyForModel("xai/grok-4")).toBeUndefined(); + }); + + test("still resolves OpenRouter multi-segment ids via first-segment split", () => { + process.env.OPENROUTER_API_KEY = "test-openrouter-key"; + expect(getApiKeyForModel("openrouter/stealth/ox-alpha")).toBe("test-openrouter-key"); + }); + }); + + describe("getModelValue", () => { + test("builds xai/ for an XAI provider row", () => { + expect( + getModelValue({ + id: "1", + name: "grok-4", + provider: "XAI", + providerLabel: null, + isPlanDefault: false, + isTaskDefault: false, + }), + ).toBe("xai/grok-4"); + }); + + test("builds anthropic/ for an ANTHROPIC provider row (regression)", () => { + expect( + getModelValue({ + id: "2", + name: "claude-sonnet-4-6", + provider: "ANTHROPIC", + providerLabel: null, + isPlanDefault: false, + isTaskDefault: false, + }), + ).toBe("anthropic/claude-sonnet-4-6"); + }); + + test("uses providerLabel prefix for OTHER rows (regression)", () => { + expect( + getModelValue({ + id: "3", + name: "stealth/ox-alpha", + provider: "OTHER", + providerLabel: "OpenRouter", + isPlanDefault: false, + isTaskDefault: false, + }), + ).toBe("openrouter/stealth/ox-alpha"); + }); }); describe("getStoredPlanModelPreference", () => { diff --git a/src/__tests__/unit/lib/ai/runCanvasAgent-xai.test.ts b/src/__tests__/unit/lib/ai/runCanvasAgent-xai.test.ts new file mode 100644 index 0000000000..fe91140c08 --- /dev/null +++ b/src/__tests__/unit/lib/ai/runCanvasAgent-xai.test.ts @@ -0,0 +1,117 @@ +/** + * Unit test: runCanvasAgent must never silently answer an xai/* model + * selection as Anthropic. aieo (^0.1.34, pinned) has no "xai" provider + * entry, so this path throws an explicit, user-visible error instead. + */ +import { describe, it, expect, vi } from "vitest"; + +vi.mock("@/lib/db", () => ({ db: {} })); +vi.mock("@/lib/pusher", () => ({ + pusherServer: { trigger: vi.fn() }, + getWorkspaceChannelName: vi.fn(() => "ch"), + PUSHER_EVENTS: { HIGHLIGHT_NODES: "highlight" }, +})); +vi.mock("@/lib/ai/askTools", () => ({ + askTools: vi.fn(() => ({})), + listConcepts: vi.fn(async () => ({ concepts: [] })), + createHasEndMarkerCondition: vi.fn(() => () => false), +})); +vi.mock("@/lib/ai/askToolsMulti", () => ({ askToolsMulti: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/workspaceConfig", () => ({ + buildWorkspaceConfigs: vi.fn(async () => [ + { + workspaceId: "ws-1", + userId: "user-1", + slug: "ws-slug", + swarmUrl: "https://swarm", + swarmApiKey: "key", + repoUrls: [], + pat: "pat", + description: "", + members: [], + currentUserGithubUsername: null, + }, + ]), + buildPublicWorkspaceConfig: vi.fn(), + fetchConceptsForWorkspaces: vi.fn(async () => ({})), +})); +vi.mock("@/lib/ai/connectionTools", () => ({ buildConnectionTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/canvasTools", () => ({ buildCanvasTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/initiativeTools", () => ({ buildInitiativeTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/researchTools", () => ({ buildResearchTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/htmlArtifactTools", () => ({ buildHtmlArtifactTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/infraTools", () => ({ buildInfraTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/graphWalkerTools", () => ({ buildGraphWalkerTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/graphWalkDispatchTools", () => ({ buildGraphWalkDispatchTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/workflowExplorerTools", () => ({ buildWorkflowExplorerTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/promptTools", () => ({ buildPromptTools: vi.fn(() => ({})) })); +vi.mock("@/lib/ai/conceptTools", () => ({ buildConceptTools: vi.fn(() => ({})) })); +vi.mock("@/lib/canvas/linkedWorkspaces", () => ({ + getLinkedWorkspacesForInitiative: vi.fn(() => []), +})); +vi.mock("@/lib/ai/message-sanitizer", () => ({ + sanitizeAndCompleteToolCalls: vi.fn(async (msgs: unknown) => msgs), +})); +vi.mock("@/lib/ai/provider", () => ({ + getModel: vi.fn(() => ({ modelId: "mock-model" })), + getApiKeyForProvider: vi.fn(() => "api-key"), +})); +vi.mock("aieo", () => ({ + getProviderOptions: vi.fn(() => ({})), + hasApiKeyForProvider: vi.fn(() => true), + PROVIDERS: ["anthropic", "google", "openai", "openrouter"], +})); +vi.mock("@/services/bifrost/orchestrator", () => ({ + getBifrostForLLM: vi.fn(async () => undefined), +})); +vi.mock("@/lib/ai/canvas-system-prompt", () => ({ + getCanvasSystemPrompt: vi.fn(async () => ({ value: "system", promptId: null })), +})); +vi.mock("@/lib/ai/capabilityGates", () => ({ + isPromptsCapabilityEnabledForOrg: vi.fn(async () => false), + isGraphWriteCapabilityEnabledForOrg: vi.fn(async () => false), + isCodeChangeCapabilityEnabledForOrg: vi.fn(async () => false), +})); +vi.mock("@/lib/constants/prompt", () => ({ + getMultiWorkspacePrefixMessages: vi.fn(() => []), + getQuickAskPrefixMessages: vi.fn(() => []), + buildCanvasScopeMessage: vi.fn(() => null), + getRoadmapCapabilitySnippet: vi.fn(() => ""), + getWhiteboardCapabilitySnippet: vi.fn(() => ""), + getPlannerCapabilitySnippet: vi.fn(() => ""), + getResearchCapabilitySnippet: vi.fn(() => ""), + getConnectionsCapabilitySnippet: vi.fn(() => ""), + getHtmlPagesCapabilitySnippet: vi.fn(() => ""), + getGraphWalkerCapabilitySnippet: vi.fn(() => ""), + getInfraCapabilitySnippet: vi.fn(() => ""), + getWorkflowsCapabilitySnippet: vi.fn(() => ""), + getPromptsCapabilitySnippet: vi.fn(() => ""), + getConceptsCapabilitySnippet: vi.fn(() => ""), + getCanvasPromptSuffix: vi.fn(() => ""), +})); + +const mockStreamText = vi.fn(); +vi.mock("ai", () => ({ + streamText: (...args: unknown[]) => mockStreamText(...args), + tool: vi.fn((t: unknown) => t), +})); + +import { runCanvasAgent } from "@/lib/ai/runCanvasAgent"; +import type { ModelMessage } from "ai"; + +describe("runCanvasAgent — xAI model selection", () => { + it("throws an explicit error for an xai/* modelName instead of silently falling back to anthropic", async () => { + await expect( + runCanvasAgent({ + userId: "user-1", + workspaceSlugs: ["ws-slug"], + messages: [{ role: "user", content: "hello" }] as ModelMessage[], + modelName: "xai/grok-4", + }), + ).rejects.toThrow(/xAI\/Grok is not yet supported/i); + + // The streamText call (and therefore any actual LLM invocation) + // must never happen for a rejected xai/* selection. + expect(mockStreamText).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/admin/llm-models/LlmModelsTable.tsx b/src/app/admin/llm-models/LlmModelsTable.tsx index 248c24f093..90dd52d31c 100644 --- a/src/app/admin/llm-models/LlmModelsTable.tsx +++ b/src/app/admin/llm-models/LlmModelsTable.tsx @@ -23,13 +23,14 @@ import { import { Switch } from "@/components/ui/switch"; import type { LlmModel, LlmProvider } from "@prisma/client"; -const PROVIDERS: LlmProvider[] = ["GOOGLE", "ANTHROPIC", "OPENAI", "AWS_BEDROCK", "OTHER"]; +const PROVIDERS: LlmProvider[] = ["GOOGLE", "ANTHROPIC", "OPENAI", "AWS_BEDROCK", "XAI", "OTHER"]; const PROVIDER_LABELS: Record = { GOOGLE: "Google", ANTHROPIC: "Anthropic", OPENAI: "OpenAI", AWS_BEDROCK: "AWS Bedrock", + XAI: "xAI", OTHER: "Other", }; diff --git a/src/app/api/admin/llm-models/[id]/route.ts b/src/app/api/admin/llm-models/[id]/route.ts index b00556b784..eb9a7bcb85 100644 --- a/src/app/api/admin/llm-models/[id]/route.ts +++ b/src/app/api/admin/llm-models/[id]/route.ts @@ -2,6 +2,63 @@ import { NextRequest, NextResponse } from "next/server"; import { requireSuperAdmin } from "@/lib/auth/require-superadmin"; import { validateApiToken } from "@/lib/auth/api-token"; import { db } from "@/lib/db"; +import { Prisma } from "@prisma/client"; + +/** + * `name` and `providerLabel` are string-concatenated into the + * `provider/name` value shipped as `vars.model` (see `getModelValue()` + * in `src/lib/ai/models.ts`), and `getApiKeyForModel` derives the + * credential from the *first* path segment. An unconstrained `name` + * containing a `/` could make a row declared as one provider resolve + * to a different provider's key. Kept in sync with the sibling + * validator in `../route.ts`. + */ +const SAFE_NAME_RE = /^[A-Za-z0-9._:-]+$/; + +function validateNameFields( + name: unknown, + providerLabel: unknown, +): NextResponse | null { + if (typeof name === "string" && !SAFE_NAME_RE.test(name)) { + return NextResponse.json( + { error: "name must match ^[A-Za-z0-9._:-]+$ (no slashes)" }, + { status: 400 }, + ); + } + if ( + typeof providerLabel === "string" && + providerLabel.length > 0 && + !SAFE_NAME_RE.test(providerLabel) + ) { + return NextResponse.json( + { error: "providerLabel must match ^[A-Za-z0-9._:-]+$ (no slashes)" }, + { status: 400 }, + ); + } + return null; +} + +/** + * Fields the shared static `API_TOKEN` (external sync services) may + * write via this endpoint. `isPublic` / `isPlanDefault` / + * `isTaskDefault` are deliberately excluded — they gate what users see + * and which model gets selected by default across the whole product, + * so a token-only caller flipping them would escalate the sync + * credential into product-wide model and spend control. Those three + * require an authenticated `SUPER_ADMIN` session (see the `isSync` + * branch below). + */ +const TOKEN_ALLOWED_FIELDS = new Set([ + "name", + "provider", + "providerLabel", + "inputPricePer1M", + "outputPricePer1M", + "cacheReadPer1MToken", + "cacheWritePer1MToken", + "dateStart", + "dateEnd", +]); export async function PATCH( request: NextRequest, @@ -22,35 +79,64 @@ export async function PATCH( } const body = await request.json(); - const { name, provider, providerLabel, inputPricePer1M, outputPricePer1M, cacheReadPer1MToken, cacheWritePer1MToken, dateStart, dateEnd, isPlanDefault, isTaskDefault, isPublic } = body; - if (isPlanDefault) { - await db.llmModel.updateMany({ where: { isPlanDefault: true, id: { not: id } }, data: { isPlanDefault: false } }); - } - if (isTaskDefault) { - await db.llmModel.updateMany({ where: { isTaskDefault: true, id: { not: id } }, data: { isTaskDefault: false } }); + if (isSync) { + const disallowed = Object.keys(body).filter( + (key) => !TOKEN_ALLOWED_FIELDS.has(key), + ); + if (disallowed.length > 0) { + return NextResponse.json( + { + error: `Token-authenticated requests cannot set: ${disallowed.join(", ")}. These require a SUPER_ADMIN session.`, + }, + { status: 403 }, + ); + } } - const model = await db.llmModel.update({ - where: { id }, - data: { - ...(name !== undefined && { name }), - ...(provider !== undefined && { provider }), - ...(providerLabel !== undefined && { providerLabel }), - ...(inputPricePer1M !== undefined && { inputPricePer1M: Number(inputPricePer1M) }), - ...(outputPricePer1M !== undefined && { outputPricePer1M: Number(outputPricePer1M) }), - ...(cacheReadPer1MToken !== undefined && { cacheReadPer1MToken: cacheReadPer1MToken != null ? Number(cacheReadPer1MToken) : null }), - ...(cacheWritePer1MToken !== undefined && { cacheWritePer1MToken: cacheWritePer1MToken != null ? Number(cacheWritePer1MToken) : null }), - ...(dateStart !== undefined && { dateStart: dateStart ? new Date(dateStart) : null }), - ...(dateEnd !== undefined && { dateEnd: dateEnd ? new Date(dateEnd) : null }), - ...(isPlanDefault !== undefined && { isPlanDefault }), - ...(isTaskDefault !== undefined && { isTaskDefault }), - ...(isPublic !== undefined && { isPublic }), - }, + const { name, provider, providerLabel, inputPricePer1M, outputPricePer1M, cacheReadPer1MToken, cacheWritePer1MToken, dateStart, dateEnd, isPlanDefault, isTaskDefault, isPublic } = body; + + const nameErr = validateNameFields(name, providerLabel); + if (nameErr) return nameErr; + + // Atomic: clearing the existing default and applying this row's + // update must happen together, or a mid-way failure can leave two + // defaults (or zero) — `getDefaultModel` resolves via `findFirst`. + const model = await db.$transaction(async (tx) => { + if (isPlanDefault) { + await tx.llmModel.updateMany({ where: { isPlanDefault: true, id: { not: id } }, data: { isPlanDefault: false } }); + } + if (isTaskDefault) { + await tx.llmModel.updateMany({ where: { isTaskDefault: true, id: { not: id } }, data: { isTaskDefault: false } }); + } + + return tx.llmModel.update({ + where: { id }, + data: { + ...(name !== undefined && { name }), + ...(provider !== undefined && { provider }), + ...(providerLabel !== undefined && { providerLabel }), + ...(inputPricePer1M !== undefined && { inputPricePer1M: Number(inputPricePer1M) }), + ...(outputPricePer1M !== undefined && { outputPricePer1M: Number(outputPricePer1M) }), + ...(cacheReadPer1MToken !== undefined && { cacheReadPer1MToken: cacheReadPer1MToken != null ? Number(cacheReadPer1MToken) : null }), + ...(cacheWritePer1MToken !== undefined && { cacheWritePer1MToken: cacheWritePer1MToken != null ? Number(cacheWritePer1MToken) : null }), + ...(dateStart !== undefined && { dateStart: dateStart ? new Date(dateStart) : null }), + ...(dateEnd !== undefined && { dateEnd: dateEnd ? new Date(dateEnd) : null }), + ...(isPlanDefault !== undefined && { isPlanDefault }), + ...(isTaskDefault !== undefined && { isTaskDefault }), + ...(isPublic !== undefined && { isPublic }), + }, + }); }); return NextResponse.json({ model }); } catch (error) { + if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { + return NextResponse.json( + { error: "A model with this name already exists" }, + { status: 409 }, + ); + } console.error("Error updating LLM model:", error); return NextResponse.json( { error: "Failed to update LLM model" }, diff --git a/src/app/api/admin/llm-models/route.ts b/src/app/api/admin/llm-models/route.ts index e8128fe214..bd140d8d3e 100644 --- a/src/app/api/admin/llm-models/route.ts +++ b/src/app/api/admin/llm-models/route.ts @@ -2,7 +2,42 @@ import { NextRequest, NextResponse } from "next/server"; import { requireSuperAdmin } from "@/lib/auth/require-superadmin"; import { validateApiToken } from "@/lib/auth/api-token"; import { db } from "@/lib/db"; -import { LlmProvider } from "@prisma/client"; +import { LlmProvider, Prisma } from "@prisma/client"; + +/** + * `name` and `providerLabel` are string-concatenated into the + * `provider/name` value shipped as `vars.model` (see `getModelValue()` + * in `src/lib/ai/models.ts`), and `getApiKeyForModel` derives the + * credential from the *first* path segment. An unconstrained `name` + * containing a `/` could make a row declared as one provider resolve + * to a different provider's key (e.g. an `XAI` row named + * "anthropic/claude-x"). No slashes, and only characters that make + * sense in a model id / display label. + */ +const SAFE_NAME_RE = /^[A-Za-z0-9._:-]+$/; + +function validateNameFields( + name: unknown, + providerLabel: unknown, +): NextResponse | null { + if (typeof name === "string" && !SAFE_NAME_RE.test(name)) { + return NextResponse.json( + { error: "name must match ^[A-Za-z0-9._:-]+$ (no slashes)" }, + { status: 400 }, + ); + } + if ( + typeof providerLabel === "string" && + providerLabel.length > 0 && + !SAFE_NAME_RE.test(providerLabel) + ) { + return NextResponse.json( + { error: "providerLabel must match ^[A-Za-z0-9._:-]+$ (no slashes)" }, + { status: 400 }, + ); + } + return null; +} export async function GET(request: NextRequest) { const authResult = await requireSuperAdmin(request); @@ -32,6 +67,12 @@ export async function POST(request: NextRequest) { if (authResult instanceof NextResponse) return authResult; } + // Captured outside the try block so the P2002 handler can look up the + // conflicting row without re-reading the request body (NextRequest + // bodies can only be consumed once — `request.clone()` after the body + // has already been read via `.json()` doesn't give a fresh stream). + let singleCreateName: string | undefined; + try { const body = await request.json(); @@ -47,6 +88,8 @@ export async function POST(request: NextRequest) { { status: 400 } ); } + const nameErr = validateNameFields(item.name, item.providerLabel); + if (nameErr) return nameErr; } const models = await Promise.all( @@ -70,14 +113,21 @@ export async function POST(request: NextRequest) { cacheReadPer1MToken: cacheReadPer1MToken != null ? Number(cacheReadPer1MToken) : null, cacheWritePer1MToken: cacheWritePer1MToken != null ? Number(cacheWritePer1MToken) : null, }, + // `provider` / `providerLabel` are intentionally omitted here. + // The nightly sync (src/lib/ai/llm-model-sync.ts, scheduled by + // vercel.json) round-trips through this same batch-upsert path, + // and its diff step normalizes providers to a coarse + // OPENAI|ANTHROPIC|GOOGLE|AWS_BEDROCK|OTHER bucket — re-applying + // that here would silently revert any row an admin has + // reclassified (e.g. XAI) within a day. The sync owns pricing; + // the admin UI is authoritative for provider classification. + // Mirrors the existing deliberate omission of isPlanDefault / + // isTaskDefault / isPublic / dateStart / dateEnd below. update: { - provider, - providerLabel: providerLabel ?? null, inputPricePer1M: Number(inputPricePer1M), outputPricePer1M: Number(outputPricePer1M), cacheReadPer1MToken: cacheReadPer1MToken != null ? Number(cacheReadPer1MToken) : null, cacheWritePer1MToken: cacheWritePer1MToken != null ? Number(cacheWritePer1MToken) : null, - // intentionally omits isPlanDefault, isTaskDefault, isPublic, dateStart, dateEnd }, }) ) @@ -88,6 +138,7 @@ export async function POST(request: NextRequest) { // Single model create path (existing behaviour) const { name, provider, providerLabel, inputPricePer1M, outputPricePer1M, cacheReadPer1MToken, cacheWritePer1MToken, dateStart, dateEnd, isPlanDefault, isTaskDefault, isPublic } = body; + singleCreateName = name; if (!name || !provider || inputPricePer1M == null || outputPricePer1M == null) { return NextResponse.json( @@ -96,32 +147,60 @@ export async function POST(request: NextRequest) { ); } - if (isPlanDefault) { - await db.llmModel.updateMany({ where: { isPlanDefault: true }, data: { isPlanDefault: false } }); - } - if (isTaskDefault) { - await db.llmModel.updateMany({ where: { isTaskDefault: true }, data: { isTaskDefault: false } }); - } + const nameErr = validateNameFields(name, providerLabel); + if (nameErr) return nameErr; + + // The default-flip (clear the existing default, then create the new + // row as the default) must be atomic — a non-transactional + // read-modify-write here can leave two defaults, or (on a mid-way + // failure) zero, and `getDefaultModel` resolves via `findFirst`. + const model = await db.$transaction(async (tx) => { + if (isPlanDefault) { + await tx.llmModel.updateMany({ where: { isPlanDefault: true }, data: { isPlanDefault: false } }); + } + if (isTaskDefault) { + await tx.llmModel.updateMany({ where: { isTaskDefault: true }, data: { isTaskDefault: false } }); + } - const model = await db.llmModel.create({ - data: { - name, - provider, - providerLabel: providerLabel ?? null, - inputPricePer1M: Number(inputPricePer1M), - outputPricePer1M: Number(outputPricePer1M), - cacheReadPer1MToken: cacheReadPer1MToken != null ? Number(cacheReadPer1MToken) : null, - cacheWritePer1MToken: cacheWritePer1MToken != null ? Number(cacheWritePer1MToken) : null, - dateStart: dateStart ? new Date(dateStart) : null, - dateEnd: dateEnd ? new Date(dateEnd) : null, - isPlanDefault: isPlanDefault ?? false, - isTaskDefault: isTaskDefault ?? false, - isPublic: isPublic ?? false, - }, + return tx.llmModel.create({ + data: { + name, + provider, + providerLabel: providerLabel ?? null, + inputPricePer1M: Number(inputPricePer1M), + outputPricePer1M: Number(outputPricePer1M), + cacheReadPer1MToken: cacheReadPer1MToken != null ? Number(cacheReadPer1MToken) : null, + cacheWritePer1MToken: cacheWritePer1MToken != null ? Number(cacheWritePer1MToken) : null, + dateStart: dateStart ? new Date(dateStart) : null, + dateEnd: dateEnd ? new Date(dateEnd) : null, + isPlanDefault: isPlanDefault ?? false, + isTaskDefault: isTaskDefault ?? false, + isPublic: isPublic ?? false, + }, + }); }); return NextResponse.json({ model }, { status: 201 }); } catch (error) { + // `name` has a DB-level unique constraint — a duplicate throws + // Prisma's P2002 rather than validating cleanly. Surface a 409 with + // the existing row's id so the admin can PATCH it directly (e.g. + // migrating a pre-existing `grok-*` row from OTHER/OpenRouter to + // XAI) instead of hitting a raw 500. `singleCreateName` was captured + // before the failed create — the request body stream has already + // been consumed by `request.json()` above and can't be re-read. + if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { + const existing = singleCreateName + ? await db.llmModel.findUnique({ where: { name: singleCreateName }, select: { id: true } }) + : null; + return NextResponse.json( + { + error: "A model with this name already exists", + existingId: existing?.id, + }, + { status: 409 }, + ); + } console.error("Error creating LLM model:", error); return NextResponse.json( { error: "Failed to create LLM model" }, diff --git a/src/app/api/agent/route.ts b/src/app/api/agent/route.ts index 7c951a3872..1999916c72 100644 --- a/src/app/api/agent/route.ts +++ b/src/app/api/agent/route.ts @@ -62,9 +62,9 @@ import { authOptions } from "@/lib/auth/nextauth"; import { getServerSession } from "next-auth/next"; import { db } from "@/lib/db"; import { EncryptionService } from "@/lib/encryption"; -import { ChatRole, ChatStatus, ArtifactType } from "@prisma/client"; +import { ChatRole, ChatStatus, ArtifactType, LlmProvider } from "@prisma/client"; import { createWebhookToken, generateWebhookSecret } from "@/lib/auth/agent-jwt"; -import { isValidModel, getApiKeyForModel } from "@/lib/ai/models"; +import { isValidModel, getApiKeyForModel, PROVIDER_API_KEY_ENV_VARS } from "@/lib/ai/models"; import { canAccessServerFeature, FEATURE_FLAGS } from "@/lib/feature-flags"; import { claimPodAndGetFrontend, updatePodRepositories, POD_PORTS, releasePodById } from "@/lib/pods"; // Deep import — see comment in services/task-workflow.ts. @@ -342,8 +342,35 @@ async function createAgentSession( headers["Authorization"] = `Bearer ${agentPassword}`; } - // Determine API key based on model + // Determine API key based on model. The Anthropic fallback only fires + // when `effectiveModel` is absent — a resolved-but-keyless model (e.g. + // xai/* with no XAI_API_KEY set) yields `undefined` here, not a + // fallback key, so the session below can end up key-less. Log that + // case (env-var name + boolean only, never the value) so it's + // diagnosable without leaking secret material. const apiKey = effectiveModel ? getApiKeyForModel(effectiveModel) : process.env.ANTHROPIC_API_KEY; + if (effectiveModel?.includes("/") && !apiKey) { + const provider = effectiveModel.split("/")[0].toUpperCase(); + const envVar = PROVIDER_API_KEY_ENV_VARS[provider]; + if (envVar) { + console.error("[Agent] model provider key missing", { + taskId, + effectiveModel, + envVar, + envVarSet: Boolean(process.env[envVar]), + }); + } + } + + // xAI bypass: the Bifrost VK reconciler falls back to anthropic's + // provider suffix for any model prefix it doesn't recognize, and its + // provider allow-list doesn't include xai — routing an xai/* session + // through Bifrost today would mint a VK pointed at the wrong + // provider. Skip the call for xai/* and use the direct XAI_API_KEY + // resolved above instead. Remove once an `aieo` gateway path + a + // Bifrost-side xai provider config both exist (see + // src/services/task-workflow.ts for the matching bypass). + const isXaiModel = effectiveModel?.startsWith("xai/") ?? false; // Bifrost routing for the goose-side LLM calls. When the rollout // flag covers this workspace, mint a per-session VK + macaroon and @@ -351,13 +378,15 @@ async function createAgentSession( // body. The agent forwards them onto every LLM call so the spend // shows up on `logs.db` as `agent-name=coding-agent`. When the flag // is off, falls back to the model-resolved key (unchanged). - const bifrost = await getBifrostForLLM(bifrostAuth, { - agentName: "coding-agent", - // Pass the selected model so the Bifrost VK reconciler resolves - // the correct provider suffix on `baseUrl` (e.g. `/genai/v1beta` - // for google/* models). Without this it defaults to anthropic. - model: effectiveModel, - }); + const bifrost = isXaiModel + ? undefined + : await getBifrostForLLM(bifrostAuth, { + agentName: "coding-agent", + // Pass the selected model so the Bifrost VK reconciler resolves + // the correct provider suffix on `baseUrl` (e.g. `/genai/v1beta` + // for google/* models). Without this it defaults to anthropic. + model: effectiveModel, + }); const sessionPayload: Record = { sessionId: taskId, @@ -434,8 +463,38 @@ export async function POST(request: NextRequest) { const body = await request.json(); const { message, taskId, artifacts = [], model } = body; - // Validate model parameter if provided - const requestModel: string | undefined = isValidModel(model) ? model : undefined; + // Validate model parameter if provided. `isValidModel` only checks that + // the prefix maps to a known provider — it does NOT check catalog + // membership, so it returns true for any "provider/anything" string. + // Short aliases (sonnet, gpt, …) aren't admin-catalog rows and skip the + // DB check below; "provider/name" strings must exist as a public, + // unexpired `LlmModel` row or an authenticated caller could force Hive + // to spend a direct provider key (e.g. XAI_API_KEY) on an unregistered + // model, bypassing the isPublic/dateEnd gate that /api/llm-models enforces. + let requestModel: string | undefined = isValidModel(model) ? model : undefined; + if (requestModel?.includes("/")) { + const [prefix, ...rest] = requestModel.split("/"); + const namePart = rest.join("/"); + const isEnumProvider = prefix.toUpperCase() in LlmProvider; + const catalogMatch = await db.llmModel.findFirst({ + where: { + name: namePart, + isPublic: true, + OR: [{ dateEnd: null }, { dateEnd: { gt: new Date() } }], + ...(isEnumProvider + ? { provider: prefix.toUpperCase() as LlmProvider } + : { + provider: LlmProvider.OTHER, + providerLabel: { equals: prefix, mode: "insensitive" as const }, + }), + }, + select: { id: true }, + }); + if (!catalogMatch) { + console.warn("[Agent] rejected non-catalog model", { model: requestModel }); + requestModel = undefined; + } + } // 1. Authenticate user const session = await getServerSession(authOptions); diff --git a/src/app/api/llm-models/route.ts b/src/app/api/llm-models/route.ts index d4450cb240..a0d0674b59 100644 --- a/src/app/api/llm-models/route.ts +++ b/src/app/api/llm-models/route.ts @@ -2,6 +2,30 @@ import { NextRequest, NextResponse } from "next/server"; import { validateApiToken } from "@/lib/auth/api-token"; import { getMiddlewareContext, requireAuth } from "@/lib/middleware/utils"; import { db } from "@/lib/db"; +import { getModelValue, PROVIDER_API_KEY_ENV_VARS, type LlmModelOption } from "@/lib/ai/models"; + +/** + * A model is only actually selectable when its provider's API key is + * configured in this running environment. Every picker + * (CanvasAgentSettingsPopover, PlanStartInput, ChatInput, + * CompactTasksList) renders whatever this endpoint returns with no + * client-side availability check — clients can't read `process.env`, + * so this is the only place the check can be correct. Filtering here + * fixes every picker at once, including the pre-existing case of a + * Google (or now xAI) row showing up with no key set. + * + * A model whose provider maps to no known env var (an `OTHER` row with + * a custom `providerLabel` we don't recognize, or the bare `OTHER` + * enum with no label) has nothing to gate on — keep it, unchanged + * from prior behavior. + */ +function isProviderKeyConfigured(model: LlmModelOption): boolean { + const value = getModelValue(model); + const prefix = value.split("/")[0].toUpperCase(); + const envVar = PROVIDER_API_KEY_ENV_VARS[prefix]; + if (!envVar) return true; + return Boolean(process.env[envVar]); +} export async function GET(request: NextRequest) { // Allow either a valid API token or an authenticated session @@ -35,5 +59,7 @@ export async function GET(request: NextRequest) { orderBy: { name: "asc" }, }); - return NextResponse.json({ models }); + const availableModels = models.filter(isProviderKeyConfigured); + + return NextResponse.json({ models: availableModels }); } diff --git a/src/app/api/user/preferences/route.ts b/src/app/api/user/preferences/route.ts index e8b84d72de..a68d011a15 100644 --- a/src/app/api/user/preferences/route.ts +++ b/src/app/api/user/preferences/route.ts @@ -4,6 +4,34 @@ import { authOptions } from "@/lib/auth/nextauth"; import { db } from "@/lib/db"; import { logger } from "@/lib/logger"; import { isValidTimezone } from "@/lib/automations/schedule"; +import { getModelValue } from "@/lib/ai/models"; + +/** + * Re-validate a stored `chatAgentModel` preference against the live, + * public, unexpired catalog. A preference persisted before a provider + * cutover (e.g. a user on `openrouter/grok-*` before xAI was onboarded + * directly) has no re-validation on read otherwise, so it would keep + * routing through the old provider forever even after an admin adds + * the new one — defeating the point of the cutover. Falls back to the + * admin-configured default (`null` → inherit) when the stored value no + * longer matches any catalog row. + */ +async function revalidateChatAgentModel(stored: string | null): Promise { + if (!stored) return stored; + const [prefix, ...rest] = stored.split("/"); + if (!prefix || rest.length === 0) return null; + const namePart = rest.join("/"); + const candidates = await db.llmModel.findMany({ + where: { + name: namePart, + isPublic: true, + OR: [{ dateEnd: null }, { dateEnd: { gt: new Date() } }], + }, + select: { id: true, name: true, provider: true, providerLabel: true, isPlanDefault: true, isTaskDefault: true }, + }); + const stillValid = candidates.some((m) => getModelValue(m) === stored); + return stillValid ? stored : null; +} /** * Authenticated user's UI preferences. Currently: @@ -30,9 +58,11 @@ export async function GET() { return NextResponse.json({ error: "User not found" }, { status: 404 }); } + const chatAgentModel = await revalidateChatAgentModel(user.chatAgentModel); + return NextResponse.json({ canvasAutonomousTurns: user.canvasAutonomousTurns, - chatAgentModel: user.chatAgentModel, + chatAgentModel, timezone: user.timezone ?? "UTC", dailyRecapEnabled: user.dailyRecapEnabled, voiceLearningEnabled: user.voiceLearningEnabled, diff --git a/src/app/api/workspaces/[slug]/legal/benchmarks/run/route.ts b/src/app/api/workspaces/[slug]/legal/benchmarks/run/route.ts index 5dfd23b61e..0e4cb827be 100644 --- a/src/app/api/workspaces/[slug]/legal/benchmarks/run/route.ts +++ b/src/app/api/workspaces/[slug]/legal/benchmarks/run/route.ts @@ -219,6 +219,18 @@ export async function POST(request: NextRequest, { params }: RouteParams) { { status: 400 }, ); } + // XAI is explicitly excluded from the standard/reasoning pair: the + // judge model is Anthropic-only by this route's contract (see + // validateModel above), and pairing a single xAI apiKey with an + // Anthropic-only judge isn't supported — the run would need two + // provider credentials, not one. Revisit once paired-model key + // handling supports two providers. + if (provider === "XAI") { + return NextResponse.json( + { error: `${label} provider "XAI" is not yet supported for benchmark runs` }, + { status: 400 }, + ); + } return null; }; diff --git a/src/app/org/[githubLogin]/_components/CanvasAgentSettingsPopover.tsx b/src/app/org/[githubLogin]/_components/CanvasAgentSettingsPopover.tsx index aa4f2b6314..7c7bd13570 100644 --- a/src/app/org/[githubLogin]/_components/CanvasAgentSettingsPopover.tsx +++ b/src/app/org/[githubLogin]/_components/CanvasAgentSettingsPopover.tsx @@ -70,13 +70,22 @@ export function CanvasAgentSettingsPopover({ }; }, []); - // Load the available models for the picker. + // Load the available models for the picker. xAI/Grok rows are excluded + // here specifically (not just left to the server-side key filter) — + // `aieo` (the canvas agent's LLM SDK, ^0.1.34) has no xai provider + // entry yet, so `runCanvasAgent` throws a visible error rather than + // silently answering as Anthropic if one is selected. Hiding it from + // this picker means users never hit that error in the first place. + // Remove this filter once aieo supports xai — see the "Add xAI" + // feature notes. useEffect(() => { let cancelled = false; fetch("/api/llm-models") .then((res) => (res.ok ? res.json() : null)) .then((data) => { - if (!cancelled && data?.models) setModels(data.models); + if (!cancelled && data?.models) { + setModels((data.models as LlmModelOption[]).filter((m) => m.provider !== "XAI")); + } }) .catch(() => { /* leave empty; picker stays hidden until a retry */ diff --git a/src/components/legal/TaskDetailsModal.tsx b/src/components/legal/TaskDetailsModal.tsx index 845a2ad43e..527b495c10 100644 --- a/src/components/legal/TaskDetailsModal.tsx +++ b/src/components/legal/TaskDetailsModal.tsx @@ -153,8 +153,13 @@ export function TaskDetailsModal({ if (cancelled || !data?.models) return; // Only providers with a known API key env var can be dispatched. // OTHER models resolve through their providerLabel (e.g. OpenRouter). + // XAI is excluded here specifically: the standard/reasoning pair + // shares a single apiKey, and the judge model is Anthropic-only — + // pairing a single xAI apiKey with an Anthropic-only judge model + // isn't supported by this run route yet (see validatePairedModel + // in the run route, which rejects XAI for the same reason). const usable = (data.models as LlmModelOption[]).filter( - (m) => !!PROVIDER_API_KEY_ENV_VARS[effectiveProvider(m)], + (m) => effectiveProvider(m) !== "XAI" && !!PROVIDER_API_KEY_ENV_VARS[effectiveProvider(m)], ); setLlmModels(usable); // Initialise the judge model here, off the local `usable` array — diff --git a/src/config/env.ts b/src/config/env.ts index e3917a6086..92bb55b546 100644 --- a/src/config/env.ts +++ b/src/config/env.ts @@ -161,6 +161,14 @@ export const optionalEnvVars = { // Callers MUST go through `isPromptsCapabilityEnabledForOrgLogin(login)`. PROMPTS_CAPABILITY_ORG_LOGINS: process.env.PROMPTS_CAPABILITY_ORG_LOGINS || "stakwork", + // Direct xAI credential for Grok models (LlmProvider.XAI rows resolve to + // `xai/` and use this key). Registered here — rather than left as a + // raw `process.env` read — so it's a typed, discoverable config value like + // its sibling provider keys. Absence is a normal, expected state in any + // environment that hasn't onboarded xAI yet: `/api/llm-models` filters out + // XAI rows when this is unset (see step 8 of the xAI feature), so pickers + // stay empty rather than erroring. + XAI_API_KEY: process.env.XAI_API_KEY || "", } as const; /** diff --git a/src/lib/ai/models.ts b/src/lib/ai/models.ts index dd23975dca..d909e61d51 100644 --- a/src/lib/ai/models.ts +++ b/src/lib/ai/models.ts @@ -30,6 +30,7 @@ export const PROVIDER_DISPLAY_LABELS: Record = { GOOGLE: "Google", AWS_BEDROCK: "AWS Bedrock", OPENROUTER: "OpenRouter", + XAI: "xAI", }; // Valid model names that can be passed from frontend @@ -67,6 +68,10 @@ export const PROVIDER_API_KEY_ENV_VARS: Record = { GOOGLE: "GOOGLE_API_KEY", AWS_BEDROCK: "AWS_BEDROCK_API_KEY", OPENROUTER: "OPENROUTER_API_KEY", + // Direct xAI credential — Grok models resolve to `xai/` and use + // this key instead of routing through OpenRouter. See the "Add xAI" + // feature notes for the Bifrost/canvas-agent gaps this doesn't cover yet. + XAI: "XAI_API_KEY", OTHER: null, }; diff --git a/src/lib/ai/runCanvasAgent.ts b/src/lib/ai/runCanvasAgent.ts index 467a360ea6..b1452abf37 100644 --- a/src/lib/ai/runCanvasAgent.ts +++ b/src/lib/ai/runCanvasAgent.ts @@ -697,7 +697,22 @@ export async function runCanvasAgent( // flag is on for the primary workspace — see `getBifrostForLLM` below. let provider: Provider = "anthropic"; if (modelName?.includes("/")) { - const prefix = modelName.split("/")[0] as Provider; + const rawPrefix = modelName.split("/")[0]; + // aieo (^0.1.34, the pinned version) has no "xai" entry in `PROVIDERS` + // — it's not just an unconfigured key, the provider itself doesn't + // exist on this path yet. Falling through to the generic branch + // below would silently answer a Grok selection as Anthropic, which + // is exactly the failure mode this feature must not have. Fail + // loudly instead. `CanvasAgentSettingsPopover` excludes xai/* rows + // from the picker, so this should only fire on a stale/tampered + // `chatAgentModel` preference. Remove this guard once `aieo` (or its + // replacement) supports xai — see the "Add xAI" feature notes. + if (rawPrefix === "xai") { + throw new Error( + `runCanvasAgent: model "${modelName}" is not available on this path — xAI/Grok is not yet supported by the canvas agent's LLM SDK (aieo). Choose a different model.`, + ); + } + const prefix = rawPrefix as Provider; if (!PROVIDERS.includes(prefix)) { console.warn( `[runCanvasAgent] model "${modelName}" has unsupported provider prefix "${prefix}"; falling back to anthropic default`, diff --git a/src/services/task-workflow.ts b/src/services/task-workflow.ts index 56dbb4a298..90ba539c00 100644 --- a/src/services/task-workflow.ts +++ b/src/services/task-workflow.ts @@ -7,7 +7,7 @@ import { buildFeatureContext } from "@/services/task-coordinator"; import { EncryptionService } from "@/lib/encryption"; import { updateTaskWorkflowStatus } from "@/lib/helpers/workflow-status"; import { getStakworkTokenReference } from "@/lib/vercel/stakwork-token"; -import { getApiKeyForModel, getDefaultModel } from "@/lib/ai/models"; +import { getApiKeyForModel, getDefaultModel, PROVIDER_API_KEY_ENV_VARS } from "@/lib/ai/models"; import { fetchChatHistory } from "@/lib/helpers/chat-history"; import { isDevelopmentMode } from "@/lib/runtime"; import type { McpServerConfig } from "@/services/mcpServers"; @@ -26,6 +26,33 @@ const encryptionService = EncryptionService.getInstance(); // rolled back) instead of being killed mid-flight and stranding the task. const STAKWORK_REQUEST_TIMEOUT_MS = 30_000; +/** + * Guards the caller-controlled `webhook` override used to continue an + * existing Stakwork project (as opposed to starting a new one at + * `${STAKWORK_BASE_URL}/projects`). This value ultimately becomes the + * URL `callStakworkAPI` POSTs the Stakwork API key AND the resolved + * provider LLM key to — an unvalidated value would let a caller + * redirect that request (and both secrets) to an arbitrary host. + * + * Only same-origin-as-`STAKWORK_BASE_URL` URLs are accepted. This + * mirrors the allowlist spirit of `src/lib/run-report/url-guard.ts` + * but is deliberately simpler: `webhook` has exactly one legitimate + * destination (continuing a Stakwork project), so origin equality is + * the whole check — no bucket/region pattern matching needed. + */ +function isAllowedStakworkWebhook(webhook: string | undefined): boolean { + if (!webhook) return false; + let webhookUrl: URL; + let baseUrl: URL; + try { + webhookUrl = new URL(webhook); + baseUrl = new URL(config.STAKWORK_BASE_URL); + } catch { + return false; + } + return webhookUrl.origin === baseUrl.origin; +} + /** * Create a task and immediately trigger Stakwork workflow * This replicates the flow: POST /api/tasks -> POST /api/chat/message @@ -825,17 +852,31 @@ export async function callStakworkAPI(params: { // the orchestrator's defaults are tuned for chat turns — caller // tuning of ttlSeconds / maxCostUsd / maxSteps is intentionally // deferred to a follow-up so this initial wiring stays small. - const bifrost = await getBifrostForLLM( - { workspaceId, workspaceSlug, userId }, - { - agentName: mode === "plan_mode" ? "plan-agent" : "coding-agent", - // Pass the selected model so the Bifrost VK reconciler resolves - // the correct provider suffix on `baseUrl` (e.g. `/genai/v1beta` - // for google/* models). Without this it defaults to anthropic - // and Google/OpenAI models get routed to the wrong provider. - model: effectiveModel ?? undefined, - }, - ); + // + // xAI bypass: `reconcileBifrostVK` derives the `baseUrl` provider + // suffix from the model prefix and falls back to anthropic for any + // prefix it doesn't recognize, and `DEFAULT_PROVIDERS` doesn't list + // "xai" — so routing an `xai/*` selection through Bifrost today would + // mint a VK pointed at the wrong provider (or error). Skip the + // Bifrost call entirely for xai/* and fall through to the direct + // `vars.apiKey` (XAI_API_KEY) resolved above. Remove this bypass once + // an `aieo` version with an xAI gateway path AND a Bifrost-side xAI + // provider config both exist — until then this trades away per-agent + // cost attribution / macaroon observability for Grok runs only. + const isXaiModel = effectiveModel?.startsWith("xai/") ?? false; + const bifrost = isXaiModel + ? undefined + : await getBifrostForLLM( + { workspaceId, workspaceSlug, userId }, + { + agentName: mode === "plan_mode" ? "plan-agent" : "coding-agent", + // Pass the selected model so the Bifrost VK reconciler resolves + // the correct provider suffix on `baseUrl` (e.g. `/genai/v1beta` + // for google/* models). Without this it defaults to anthropic + // and Google/OpenAI models get routed to the wrong provider. + model: effectiveModel ?? undefined, + }, + ); if (bifrost) { vars.apiKey = bifrost.apiKey; vars.baseUrl = bifrost.baseUrl; @@ -852,16 +893,42 @@ export async function callStakworkAPI(params: { // route for this dispatch. Look for "[callStakworkAPI] model routing" // in Vercel logs (filter by /api/chat/message). `baseUrl` should carry // the model's provider suffix (e.g. /genai/v1beta for google/* models). + // + // Deliberately excludes any substring of the resolved key — a prior + // version logged `vars.apiKey.slice(0, 7)` as `apiKeyPrefix`, which + // for a short fixed-prefix key (e.g. xAI's `xai-...`) is real key + // material, not a discriminator. `providerKeySet` (derived from the + // model prefix) replaces the Google-only `googleKeySet` so a missing + // key is visible for any provider, not just Google. + const routingProvider = effectiveModel?.includes("/") ? effectiveModel.split("/")[0].toUpperCase() : null; + const routingEnvVar = routingProvider ? PROVIDER_API_KEY_ENV_VARS[routingProvider] : null; console.log("[callStakworkAPI] model routing", { taskId, mode, effectiveModel, bifrostActive: Boolean(bifrost), baseUrl: vars.baseUrl, - apiKeyPrefix: typeof vars.apiKey === "string" ? vars.apiKey.slice(0, 7) : null, - googleKeySet: Boolean(process.env.GOOGLE_API_KEY), + providerKeySet: routingEnvVar ? Boolean(process.env[routingEnvVar]) : null, }); + // A prefixed model whose provider maps to a real env var but resolved + // to no key anywhere (not from `getApiKeyForModel` above, not from + // Bifrost) means the dispatch is about to go out key-less. Log which + // env var is missing — name + boolean only, never the value — so this + // is diagnosable in Vercel logs without leaking secret material. This + // is a log-only change; control flow is unaffected and any subsequent + // `{ error }` this function returns must stay generic ("model + // provider not configured") with no env-var names in the HTTP response. + if (routingEnvVar && !vars.apiKey) { + console.error("[callStakworkAPI] model provider key missing", { + taskId, + mode, + effectiveModel, + envVar: routingEnvVar, + envVarSet: Boolean(process.env[routingEnvVar]), + }); + } + // Get workflow ID (replicating workflow selection logic) const stakworkWorkflowIds = config.STAKWORK_WORKFLOW_ID.split(","); @@ -902,8 +969,25 @@ export async function callStakworkAPI(params: { }; // Make Stakwork API call (replicating fetch call from chat/message route) - // If webhook is provided, use it to continue existing workflow; otherwise start new project - const stakworkURL = webhook || `${config.STAKWORK_BASE_URL}/projects`; + // If webhook is provided, use it to continue existing workflow; otherwise start new project. + // + // `webhook` arrives on a caller-controlled request body (see + // /api/chat/message and the roadmap feature-chat dispatcher) and this + // request carries `Authorization: Token token=${STAKWORK_API_KEY}` plus + // `vars.apiKey` (the resolved provider LLM key) in its body — an + // unvalidated `webhook` value lets any caller redirect that fetch to an + // attacker-chosen host and exfiltrate both secrets. Only accept it when + // its origin matches `STAKWORK_BASE_URL`; anything else falls back to + // the default `/projects` endpoint (silently — same behavior as if + // `webhook` had been omitted) with a warning logged for visibility. + const stakworkURL = isAllowedStakworkWebhook(webhook) + ? webhook! + : `${config.STAKWORK_BASE_URL}/projects`; + if (webhook && webhook !== stakworkURL) { + console.warn("[callStakworkAPI] rejected webhook URL with non-Stakwork origin; using default", { + taskId, + }); + } try { const response = await fetch(stakworkURL, { diff --git a/src/utils/mockSetup.ts b/src/utils/mockSetup.ts index cac1a1508b..47e7d454cb 100644 --- a/src/utils/mockSetup.ts +++ b/src/utils/mockSetup.ts @@ -1026,6 +1026,20 @@ export async function ensureMockLlmModels(): Promise { inputPricePer1M: 3, outputPricePer1M: 15, }, + { + // xAI/Grok — exercises "users pick it from the existing model + // dropdowns" for the new direct-provider path before production. + // Like the Anthropic rows above, this must be a real xAI model id + // since it's threaded straight through as the model id. Note this + // row is still subject to the `/api/llm-models` env-key filter + // (see isProviderKeyConfigured in src/app/api/llm-models/route.ts) + // — it only appears in pickers when `XAI_API_KEY` is also set in + // this environment, same as any other provider. + name: "grok-4", + provider: LlmProvider.XAI, + inputPricePer1M: 3, + outputPricePer1M: 15, + }, ]; for (const m of models) { From 88050a44b1d042c6c7207cf65d398cdf1ee0ae7b Mon Sep 17 00:00:00 2001 From: tomsmith8 Date: Wed, 2 Sep 2026 15:08:54 +0000 Subject: [PATCH 2/2] Generated with Hive: Add origin validation tests for webhook URL in callStakworkAPI --- .../unit/services/call-stakwork-api.test.ts | 19 +++++++++++- .../unit/services/task-workflow.test.ts | 29 ++++++++++++++++++- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/src/__tests__/unit/services/call-stakwork-api.test.ts b/src/__tests__/unit/services/call-stakwork-api.test.ts index 01531fd8f0..cf53806b80 100644 --- a/src/__tests__/unit/services/call-stakwork-api.test.ts +++ b/src/__tests__/unit/services/call-stakwork-api.test.ts @@ -339,7 +339,11 @@ describe("callStakworkAPI", () => { test("uses webhook URL when provided (for FORM artifact continuation)", async () => { const { config } = await import("@/config/env"); - const webhookUrl = "https://stakwork.example.com/webhook/continue/abc123"; + // Must be same-origin as STAKWORK_BASE_URL — callStakworkAPI only + // honors a caller-supplied `webhook` override when its origin + // matches, since that URL receives the Stakwork API key and the + // resolved provider LLM key (see isAllowedStakworkWebhook). + const webhookUrl = `${config.STAKWORK_BASE_URL}/webhook/continue/abc123`; mockFetch.mockResolvedValueOnce(createSuccessResponse() as any); await callStakworkAPI(createTestParams({ webhook: webhookUrl })); @@ -350,6 +354,19 @@ describe("callStakworkAPI", () => { ); }); + test("falls back to /projects endpoint when webhook has a different origin than STAKWORK_BASE_URL", async () => { + const { config } = await import("@/config/env"); + const foreignWebhookUrl = "https://attacker.example.com/webhook/continue/abc123"; + mockFetch.mockResolvedValueOnce(createSuccessResponse() as any); + + await callStakworkAPI(createTestParams({ webhook: foreignWebhookUrl })); + + expect(mockFetch).toHaveBeenCalledWith( + `${config.STAKWORK_BASE_URL}/projects`, + expect.any(Object) + ); + }); + test("falls back to /projects endpoint when webhook is not provided", async () => { const { config } = await import("@/config/env"); mockFetch.mockResolvedValueOnce(createSuccessResponse() as any); diff --git a/src/__tests__/unit/services/task-workflow.test.ts b/src/__tests__/unit/services/task-workflow.test.ts index c6e18cce07..67d2595b28 100644 --- a/src/__tests__/unit/services/task-workflow.test.ts +++ b/src/__tests__/unit/services/task-workflow.test.ts @@ -3024,7 +3024,12 @@ describe("callStakworkAPI - Direct Unit Tests", () => { json: async () => TestDataFactory.createStakworkSuccessResponse(), } as Response); - const customWebhookUrl = "https://custom-stakwork.com/api/continue-workflow"; + // Must be same-origin as the mocked STAKWORK_BASE_URL + // ("https://test-stakwork.com") — callStakworkAPI only honors a + // caller-supplied `webhook` override when its origin matches, + // since that URL receives the Stakwork API key and the resolved + // provider LLM key (see isAllowedStakworkWebhook). + const customWebhookUrl = "https://test-stakwork.com/api/continue-workflow"; const params = TestDataFactory.createCallStakworkAPIParams({ webhook: customWebhookUrl, }); @@ -3040,6 +3045,28 @@ describe("callStakworkAPI - Direct Unit Tests", () => { ); }); + test("should fall back to /projects endpoint when webhook has a different origin than STAKWORK_BASE_URL", async () => { + mockFetch.mockResolvedValue({ + ok: true, + json: async () => TestDataFactory.createStakworkSuccessResponse(), + } as Response); + + const foreignWebhookUrl = "https://attacker.example.com/api/continue-workflow"; + const params = TestDataFactory.createCallStakworkAPIParams({ + webhook: foreignWebhookUrl, + }); + + const { callStakworkAPI } = await import("@/services/task-workflow"); + await callStakworkAPI(params); + + expect(mockFetch).toHaveBeenCalledWith( + "https://test-stakwork.com/projects", + expect.objectContaining({ + method: "POST", + }) + ); + }); + test("should include correct authorization header", async () => { mockFetch.mockResolvedValue({ ok: true,