From aba51aac45c57a6e13d2f34f9e1336bf7779c21c Mon Sep 17 00:00:00 2001 From: Evanfeenstra Date: Thu, 10 Sep 2026 13:47:23 -0700 Subject: [PATCH] bifrost: route every provider through the gateway, fall back per VK grant (aieo 0.1.38) Follow-up to #5296, which granted xai on VKs but left the xai/ bypasses in place. Removes the bypasses in task-workflow and the agent route, opens the canvas agent's Anthropic-only gate and passes the resolved model through, so all five providers ride Bifrost on every call site. aieo 0.1.38 prefixes OpenRouter model ids on the gateway path, which is what Bifrost routes on. The reconciler now snapshots the VK's provider grants on WorkspaceMember.bifrostVkProviders (new column + migration) and reports whether the requested model's provider is granted. The orchestrator returns undefined when it isn't, so the caller keeps its direct key instead of failing at the gateway; unknown grants fail open as before. A miss re-checks the gateway after 15 minutes instead of the 24h window. --- package-lock.json | 8 +- package.json | 2 +- .../migration.sql | 2 + prisma/schema.prisma | 5 + .../services/bifrost/orchestrator.test.ts | 79 ++++++- .../unit/services/bifrost/reconciler.test.ts | 222 +++++++++++++++++- src/app/api/agent/route.ts | 32 +-- src/lib/ai/runCanvasAgent.ts | 59 +++-- src/services/bifrost/constants.ts | 31 +++ src/services/bifrost/orchestrator.ts | 44 +++- src/services/bifrost/reconciler.ts | 153 ++++++++++-- src/services/bifrost/types.ts | 22 ++ src/services/task-workflow.ts | 42 ++-- 13 files changed, 592 insertions(+), 109 deletions(-) create mode 100644 prisma/migrations/20260910203512_add_workspace_member_bifrost_vk_providers/migration.sql diff --git a/package-lock.json b/package-lock.json index 484603fb47..bb20c4b817 100644 --- a/package-lock.json +++ b/package-lock.json @@ -65,7 +65,7 @@ "@vercel/functions": "^3.1.4", "@xyflow/react": "^12.8.6", "ai": "^6.0.72", - "aieo": "^0.1.37", + "aieo": "^0.1.38", "autoprefixer": "^10.4.21", "axios": "1.10.0", "bcryptjs": "^3.0.2", @@ -15593,9 +15593,9 @@ } }, "node_modules/aieo": { - "version": "0.1.37", - "resolved": "https://registry.npmjs.org/aieo/-/aieo-0.1.37.tgz", - "integrity": "sha512-Nuhmhn2tV39aHIKqoXIfqsg2W6Jz71QaO3O9fhVqZKYM2ec6mrtptwGgQLQOSGV91xFhYu5wS+ahTYUccJrq9g==", + "version": "0.1.38", + "resolved": "https://registry.npmjs.org/aieo/-/aieo-0.1.38.tgz", + "integrity": "sha512-DqTIFwfnyGAi0EkxIxvgsYeUXymJhEczYeG7HKZZV7qFV1kPrGV4lVm3JDkm0JE6+Emn+08qaR0eBxjQfUrYuQ==", "license": "ISC", "dependencies": { "@ai-sdk/anthropic": "3.0.105", diff --git a/package.json b/package.json index 1798ab82b9..741be3227c 100644 --- a/package.json +++ b/package.json @@ -105,7 +105,7 @@ "@vercel/functions": "^3.1.4", "@xyflow/react": "^12.8.6", "ai": "^6.0.72", - "aieo": "^0.1.37", + "aieo": "^0.1.38", "autoprefixer": "^10.4.21", "axios": "1.10.0", "bcryptjs": "^3.0.2", diff --git a/prisma/migrations/20260910203512_add_workspace_member_bifrost_vk_providers/migration.sql b/prisma/migrations/20260910203512_add_workspace_member_bifrost_vk_providers/migration.sql new file mode 100644 index 0000000000..11d73397e2 --- /dev/null +++ b/prisma/migrations/20260910203512_add_workspace_member_bifrost_vk_providers/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "workspace_members" ADD COLUMN "bifrost_vk_providers" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index bdf747a013..da036a6408 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -358,6 +358,11 @@ model WorkspaceMember { bifrostVkId String? @map("bifrost_vk_id") bifrostCustomerId String? @map("bifrost_customer_id") bifrostSyncedAt DateTime? @map("bifrost_synced_at") + // Provider grants the cached Bifrost VK carried when it was last + // created or refreshed (`bifrost_synced_at`). Empty = unknown (row + // predates this column or the gateway couldn't be read), never "no + // grants" — the reconciler treats empty as "refresh, then trust Bifrost". + bifrostVkProviders String[] @default([]) @map("bifrost_vk_providers") user User @relation(fields: [userId], references: [id], onDelete: Cascade) workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade) diff --git a/src/__tests__/unit/services/bifrost/orchestrator.test.ts b/src/__tests__/unit/services/bifrost/orchestrator.test.ts index a6a9313799..c85e3e43e3 100644 --- a/src/__tests__/unit/services/bifrost/orchestrator.test.ts +++ b/src/__tests__/unit/services/bifrost/orchestrator.test.ts @@ -1,4 +1,5 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import type { BifrostProvider } from "@/services/bifrost/types"; // Mock the per-layer reconcilers BEFORE importing the SUT. The // orchestrator imports them by their concrete module paths @@ -68,7 +69,15 @@ function mockMintOk(runId = "run_fixed_abc") { }); } -function mockVKOk(overrides: Partial<{ vkValue: string; baseUrl: string }> = {}) { +function mockVKOk( + overrides: Partial<{ + vkValue: string; + baseUrl: string; + modelProvider: BifrostProvider; + providers: string[]; + modelProviderGranted: boolean; + }> = {}, +) { vi.mocked(reconcilerModule.reconcileBifrostVK).mockResolvedValueOnce({ workspaceId: "ws-1", userId: "u_alice", @@ -77,6 +86,9 @@ function mockVKOk(overrides: Partial<{ vkValue: string; baseUrl: string }> = {}) vkValue: overrides.vkValue ?? "sk-bf-LIVE", baseUrl: overrides.baseUrl ?? "http://bifrost.test:8181", created: false, + modelProvider: overrides.modelProvider ?? "anthropic", + providers: overrides.providers, + modelProviderGranted: overrides.modelProviderGranted, }); } @@ -201,6 +213,70 @@ describe("getBifrostForLLM (master reconciler)", () => { ); }); + // ── Grant check ───────────────────────────────────────────────────── + + it("returns undefined without minting when the VK has no grant for the model's provider", async () => { + process.env.BIFROST_ENABLED = "true"; + vi.spyOn(console, "info").mockImplementation(() => {}); + vi.spyOn(console, "log").mockImplementation(() => {}); + mockVKOk({ + baseUrl: "http://bifrost.test:8181/openai/v1", + modelProvider: "xai", + providers: ["anthropic", "openai", "openrouter", "gemini"], + modelProviderGranted: false, + }); + + const result = await getBifrostForLLM(auth, { + agentName: "coding-agent", + model: "xai/grok-4.6", + }); + // Caller keeps its direct provider key — the gateway would have + // rejected the call, and nothing downstream retries. + expect(result).toBeUndefined(); + expect(reconcilerModule.reconcileBifrostVK).toHaveBeenCalledWith( + "ws-1", + "u_alice", + { model: "xai/grok-4.6" }, + ); + expect(issuerModule.mintInvocationMacaroon).not.toHaveBeenCalled(); + }); + + it("routes through Bifrost when the grants are unknown (no snapshot yet)", async () => { + process.env.BIFROST_ENABLED = "true"; + mockVKOk({ + baseUrl: "http://bifrost.test:8181/openai/v1", + modelProvider: "xai", + providers: undefined, + modelProviderGranted: undefined, + }); + mockMintOk(); + + const result = await getBifrostForLLM(auth, { + agentName: "coding-agent", + model: "xai/grok-4.6", + }); + expect(result?.baseUrl).toBe("http://bifrost.test:8181/openai/v1"); + expect(issuerModule.mintInvocationMacaroon).toHaveBeenCalledTimes(1); + }); + + it("routes through Bifrost when the model's provider is granted", async () => { + process.env.BIFROST_ENABLED = "true"; + mockVKOk({ + baseUrl: "http://bifrost.test:8181/openai/v1", + modelProvider: "xai", + providers: ["anthropic", "openai", "openrouter", "gemini", "xai"], + modelProviderGranted: true, + }); + mockMintOk(); + + const result = await getBifrostForLLM(auth, { + agentName: "coding-agent", + model: "xai/grok-4.6", + }); + expect(result?.apiKey).toBe("sk-bf-LIVE"); + expect(issuerModule.mintInvocationMacaroon).toHaveBeenCalledTimes(1); + }); + it("forwards caller-supplied runId + budget overrides to the issuer", async () => { process.env.BIFROST_ENABLED = "true"; mockVKOk(); @@ -396,6 +472,7 @@ describe("getBifrostForLLM (master reconciler)", () => { vkValue: "sk-bf-LIVE", baseUrl: "http://bifrost.test:8181/anthropic/v1", created: false, + modelProvider: "anthropic", }; }, ); diff --git a/src/__tests__/unit/services/bifrost/reconciler.test.ts b/src/__tests__/unit/services/bifrost/reconciler.test.ts index 09a171fb19..8bdf67f8e1 100644 --- a/src/__tests__/unit/services/bifrost/reconciler.test.ts +++ b/src/__tests__/unit/services/bifrost/reconciler.test.ts @@ -147,6 +147,9 @@ describe("reconcileBifrostVK", () => { // LLM-facing URL we hand back to callers. baseUrl: "http://bifrost.test:8181/anthropic/v1", created: false, + // Same default: no model → anthropic. No grant snapshot on the + // row and a fresh stamp → grants unknown, fail open. + modelProvider: "anthropic", }); expect(client.listCustomers).not.toHaveBeenCalled(); expect(client.listProviders).not.toHaveBeenCalled(); @@ -794,7 +797,10 @@ describe("provider grants", () => { bifrostCustomerId: null, } as never); } - function cachedMember(bifrostSyncedAt: Date | null) { + function cachedMember( + bifrostSyncedAt: Date | null, + bifrostVkProviders: string[] = [], + ) { vi.mocked(dbMock.workspaceMember.findUnique).mockResolvedValueOnce({ id: "mem-1", bifrostVkValue: JSON.stringify({ @@ -807,6 +813,7 @@ describe("provider grants", () => { bifrostVkId: "vk-1", bifrostCustomerId: "cust-1", bifrostSyncedAt, + bifrostVkProviders, } as never); } /** Client for the create path: customer exists, no VK yet. */ @@ -1098,7 +1105,10 @@ describe("provider grants", () => { ).toEqual(ALL_PROVIDERS); expect(dbMock.workspaceMember.update).toHaveBeenCalledWith({ where: { id: "mem-1" }, - data: { bifrostSyncedAt: expect.any(Date) }, + data: { + bifrostSyncedAt: expect.any(Date), + bifrostVkProviders: ALL_PROVIDERS, + }, }); }); @@ -1118,7 +1128,10 @@ describe("provider grants", () => { expect(client.updateVirtualKey).not.toHaveBeenCalled(); expect(dbMock.workspaceMember.update).toHaveBeenCalledWith({ where: { id: "mem-1" }, - data: { bifrostSyncedAt: expect.any(Date) }, + data: { + bifrostSyncedAt: expect.any(Date), + bifrostVkProviders: ALL_PROVIDERS, + }, }); }); @@ -1144,4 +1157,207 @@ describe("provider grants", () => { }); }); }); + + describe("grant snapshot → model provider", () => { + const MIN_MS = 60 * 1000; + + it("reports the Bifrost provider the model maps to", async () => { + const cases: Array<[string | undefined, string]> = [ + [undefined, "anthropic"], + ["anthropic/claude-sonnet-4-6", "anthropic"], + ["google/gemini-2.5-flash", "gemini"], + ["xai/grok-4.6", "xai"], + ["openrouter/moonshotai/kimi-k2", "openrouter"], + ["gpt", "openai"], + ]; + for (const [model, expected] of cases) { + cachedMember(new Date(), ALL_PROVIDERS); + const client = makeClientStub(); + const result = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => client, + model, + }); + expect(result.modelProvider).toBe(expected); + expect(result.providers).toEqual(ALL_PROVIDERS); + expect(result.modelProviderGranted).toBe(true); + expect(client.getVirtualKey).not.toHaveBeenCalled(); + } + }); + + it("flags a model whose provider the fresh snapshot lacks, without touching the gateway", async () => { + cachedMember(new Date(), FOUR); + const client = makeClientStub(); + + const result = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => client, + model: "xai/grok-4.6", + }); + + expect(result.vkValue).toBe("sk-bf-CACHED"); + expect(result.baseUrl).toBe("http://bifrost.test:8181/openai/v1"); + expect(result.modelProvider).toBe("xai"); + expect(result.providers).toEqual(FOUR); + expect(result.modelProviderGranted).toBe(false); + expect(client.listProviders).not.toHaveBeenCalled(); + expect(client.getVirtualKey).not.toHaveBeenCalled(); + expect(dbMock.workspaceMember.update).not.toHaveBeenCalled(); + }); + + it("re-checks after the miss window and grants a provider the gateway now has", async () => { + // Past the 15-minute miss window, well inside the 24h one. + cachedMember(new Date(Date.now() - 20 * MIN_MS), FOUR); + const updateVirtualKey = vi.fn().mockResolvedValue({}); + const client = makeClientStub({ + getVirtualKey: vi + .fn() + .mockResolvedValue({ virtual_key: existingVk(fourConfigs()) }), + updateVirtualKey, + }); + + const result = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => client, + model: "xai/grok-4.6", + }); + + expect(client.getVirtualKey).toHaveBeenCalledWith("vk-1"); + expect(updateVirtualKey).toHaveBeenCalledTimes(1); + expect(result.providers).toEqual(ALL_PROVIDERS); + expect(result.modelProviderGranted).toBe(true); + expect(dbMock.workspaceMember.update).toHaveBeenCalledWith({ + where: { id: "mem-1" }, + data: { + bifrostSyncedAt: expect.any(Date), + bifrostVkProviders: ALL_PROVIDERS, + }, + }); + }); + + it("keeps falling back while the gateway lacks the provider, one re-check per miss window", async () => { + cachedMember(new Date(Date.now() - 20 * MIN_MS), FOUR); + const client = makeClientStub({ + listProviders: vi.fn().mockResolvedValue(providersResponse(FOUR)), + getVirtualKey: vi + .fn() + .mockResolvedValue({ virtual_key: existingVk(fourConfigs()) }), + }); + + const result = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => client, + model: "xai/grok-4.6", + }); + + expect(client.getVirtualKey).toHaveBeenCalledTimes(1); + expect(client.updateVirtualKey).not.toHaveBeenCalled(); + expect(result.providers).toEqual(FOUR); + expect(result.modelProviderGranted).toBe(false); + // Re-stamped, so the next call inside the window skips the gateway. + expect(dbMock.workspaceMember.update).toHaveBeenCalledWith({ + where: { id: "mem-1" }, + data: { + bifrostSyncedAt: expect.any(Date), + bifrostVkProviders: FOUR, + }, + }); + }); + + it("treats an empty snapshot as unknown: fails open inside the miss window, refreshes after it", async () => { + // A row from before the snapshot column, stamped a minute ago. + cachedMember(new Date(Date.now() - MIN_MS)); + const early = makeClientStub(); + const first = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => early, + model: "xai/grok-4.6", + }); + expect(first.providers).toBeUndefined(); + expect(first.modelProviderGranted).toBeUndefined(); + expect(early.getVirtualKey).not.toHaveBeenCalled(); + + cachedMember(new Date(Date.now() - 20 * MIN_MS)); + const later = makeClientStub({ + getVirtualKey: vi.fn().mockResolvedValue({ + virtual_key: existingVk([...fourConfigs(), hydratedConfig(5, "xai")]), + }), + }); + const second = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => later, + model: "xai/grok-4.6", + }); + expect(later.getVirtualKey).toHaveBeenCalledWith("vk-1"); + expect(later.updateVirtualKey).not.toHaveBeenCalled(); + expect(second.providers).toEqual(ALL_PROVIDERS); + expect(second.modelProviderGranted).toBe(true); + }); + + it("keeps the stale snapshot (and falls back) when the refresh itself fails", async () => { + cachedMember(new Date(Date.now() - 20 * MIN_MS), FOUR); + const client = makeClientStub({ + getVirtualKey: vi + .fn() + .mockRejectedValue(new BifrostHttpError(500, undefined, "boom")), + }); + + const result = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => client, + model: "xai/grok-4.6", + }); + + expect(result.vkValue).toBe("sk-bf-CACHED"); + expect(result.providers).toEqual(FOUR); + expect(result.modelProviderGranted).toBe(false); + // Stamped without touching the snapshot a failed read didn't see. + expect(dbMock.workspaceMember.update).toHaveBeenCalledWith({ + where: { id: "mem-1" }, + data: { bifrostSyncedAt: expect.any(Date) }, + }); + }); + + it("snapshots the grants on create and flags a model the gateway can't serve", async () => { + uncachedMember(); + const client = creatingClient({ + listProviders: vi.fn().mockResolvedValue(providersResponse(FOUR)), + }); + + const result = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => client, + model: "xai/grok-4.6", + }); + + expect(grantedProviders(client)).toEqual(FOUR); + expect(result.modelProvider).toBe("xai"); + expect(result.providers).toEqual(FOUR); + expect(result.modelProviderGranted).toBe(false); + expect(dbMock.workspaceMember.update).toHaveBeenCalledWith( + expect.objectContaining({ + data: expect.objectContaining({ bifrostVkProviders: FOUR }), + }), + ); + }); + + it("snapshots an existing VK's grants after topping it up", async () => { + uncachedMember(); + const updateVirtualKey = vi.fn().mockResolvedValue({}); + const client = creatingClient({ + listVirtualKeys: vi + .fn() + .mockResolvedValue(vkList([existingVk(fourConfigs())])), + updateVirtualKey, + }); + + const result = await reconcileBifrostVK(WORKSPACE_ID, USER_ID, { + clientFactory: () => client, + model: "xai/grok-4.6", + }); + + expect(client.createVirtualKey).not.toHaveBeenCalled(); + expect(updateVirtualKey).toHaveBeenCalledTimes(1); + expect(result.providers).toEqual(ALL_PROVIDERS); + expect(result.modelProviderGranted).toBe(true); + expect(dbMock.workspaceMember.update).toHaveBeenCalledWith( + expect.objectContaining({ + data: expect.objectContaining({ bifrostVkProviders: ALL_PROVIDERS }), + }), + ); + }); + }); + }); diff --git a/src/app/api/agent/route.ts b/src/app/api/agent/route.ts index fd0720ea88..4ccb77220a 100644 --- a/src/app/api/agent/route.ts +++ b/src/app/api/agent/route.ts @@ -362,32 +362,22 @@ async function createAgentSession( } } - // xAI bypass: DEFAULT_PROVIDERS (src/services/bifrost/constants.ts) - // now lists xai and the VK reconciler grants it on any gateway that - // has it configured, but the swarm gateways don't yet carry an - // XAI_API_KEY (sphinx-swarm's bifrost config is the missing piece) — - // routing an xai/* session through Bifrost today would still fail on - // most swarms. Skip the call for xai/* and use the direct XAI_API_KEY - // resolved above instead. Remove once the swarm gateways ship with an - // xai key (aieo already maps xai onto the /openai/v1 gateway path). - // 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 // override `apiKey` + inject `baseUrl` / `headers` on the session // 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 = 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, - }); + // is off — or the VK has no grant for the model's provider on this + // swarm's gateway, which the orchestrator checks — falls back to the + // model-resolved key (unchanged). Every provider rides Bifrost here, + // xai included; src/services/task-workflow.ts has the same wiring. + 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 sessionPayload: Record = { sessionId: taskId, diff --git a/src/lib/ai/runCanvasAgent.ts b/src/lib/ai/runCanvasAgent.ts index f471e31550..5334c92c0b 100644 --- a/src/lib/ai/runCanvasAgent.ts +++ b/src/lib/ai/runCanvasAgent.ts @@ -322,13 +322,14 @@ export interface RunCanvasAgentOptions { * Optional model override in `getModelValue()` "provider/name" form * (e.g. "anthropic/claude-opus-4-6", "openrouter/stealth/ox-alpha"), * sourced from the caller's `User.chatAgentModel` preference. Any - * aieo-supported provider prefix (anthropic/google/openai/openrouter) - * with a configured API key is honored — the provider, key, and - * provider options all follow the prefix. Anthropic-only features - * degrade gracefully off-Anthropic: the `web_search` provider tool is - * stripped, and Bifrost routing / prompt caching don't apply. An - * unsupported prefix or missing key logs a warning and falls back to - * aieo's Anthropic default (sonnet), as does omitting this entirely. + * aieo-supported provider prefix (anthropic/google/openai/openrouter/ + * xai) with a configured API key is honored — the provider, key, + * provider options and Bifrost routing all follow the prefix. + * Anthropic-only features degrade gracefully off-Anthropic: the + * native `web_search` tool gives way to the HTTP shim, and prompt + * caching doesn't apply. An unsupported prefix or missing key logs a + * warning and falls back to aieo's Anthropic default (sonnet), as + * does omitting this entirely. */ modelName?: string; /** @@ -662,7 +663,6 @@ function primaryOwnerRepo(repoUrls: string[]): string | undefined { export async function runCanvasAgent( opts: RunCanvasAgentOptions, ): Promise { - const start = Date.now(); const { userId, orgId, @@ -716,8 +716,9 @@ export async function runCanvasAgent( // default sonnet), LOUDLY: a selected model must never silently // answer as a different one. The `getModel` call is deferred until // after workspace resolution so we can thread Bifrost overrides - // (baseUrl + `x-macaroon` headers, Anthropic-only) when the rollout - // flag is on for the primary workspace — see `getBifrostForLLM` below. + // (baseUrl + `x-macaroon` headers, for whichever provider won) when + // the rollout flag is on for the primary workspace — see + // `getBifrostForLLM` below. let provider: Provider = "anthropic"; if (modelName?.includes("/")) { const prefix = modelName.split("/")[0] as Provider; @@ -1238,18 +1239,37 @@ export async function runCanvasAgent( // mirrors the `repo-agent` vs `diagram-agent` convention of naming // by user-facing purpose, not by underlying function. const agentName = orgId ? "canvas-agent" : "chat-agent"; + + // Honor the caller's model preference when it targets the resolved + // provider (which was itself derived from the same prefix, so any + // supported selection with a configured key passes). `modelName` is + // in "provider/name" form; aieo strips the known-provider prefix and + // uses the remainder as the model id (multi-segment ids like + // "openrouter/stealth/ox-alpha" → "stealth/ox-alpha"). A selection + // that failed provider resolution above stays undefined here so we + // never pair a foreign model id with the fallback Anthropic client. + const modelOverride = + modelName && modelName.startsWith(`${provider}/`) ? modelName : undefined; + const tBifrost = Date.now(); - // Bifrost VKs and gateway routing are Anthropic-only; a non-Anthropic - // model preference goes straight to its own provider endpoint. + // Every provider rides Bifrost. The VK reconciler suffixes `baseUrl` + // for the model's provider (`/openai/v1` for openai, openrouter and + // xai; `/genai/v1beta` for google) and aieo prefixes the model id on + // the OpenAI-compat route so the gateway picks the right upstream. + // The model handed over is the one `getModel` will actually use — + // undefined means aieo's Anthropic default. When the VK has no grant + // for that provider on this workspace's gateway, the orchestrator + // returns undefined and the call goes straight to the provider with + // Hive's own key, exactly as before. const bifrost = - provider === "anthropic" && primaryWorkspaceId && primaryUserId + primaryWorkspaceId && primaryUserId ? await getBifrostForLLM( { workspaceId: primaryWorkspaceId, workspaceSlug: primarySlug, userId: primaryUserId, }, - { agentName }, + { agentName, model: modelOverride }, ) : undefined; const tBifrostMs = Date.now() - tBifrost; @@ -1259,17 +1279,6 @@ export async function runCanvasAgent( console.log("[runCanvasAgent] timing", { stage: "getBifrostForLLM", ms: 0, skipped: "no primaryWorkspaceId/userId", workspaces: workspaceSlugs, orgId: orgId ?? null }); } - // Honor the caller's model preference when it targets the resolved - // provider (which was itself derived from the same prefix, so any - // supported selection with a configured key passes). `modelName` is - // in "provider/name" form; aieo strips the known-provider prefix and - // uses the remainder as the model id (multi-segment ids like - // "openrouter/stealth/ox-alpha" → "stealth/ox-alpha"). A selection - // that failed provider resolution above stays undefined here so we - // never pair a foreign model id with the fallback Anthropic client. - const modelOverride = - modelName && modelName.startsWith(`${provider}/`) ? modelName : undefined; - const model = getModel( provider, bifrost?.apiKey ?? apiKey, diff --git a/src/services/bifrost/constants.ts b/src/services/bifrost/constants.ts index 4e8da5efc6..96d84081d9 100644 --- a/src/services/bifrost/constants.ts +++ b/src/services/bifrost/constants.ts @@ -1,3 +1,5 @@ +import type { Provider as AieoProvider } from "aieo"; + import type { BifrostProvider } from "./types"; /** @@ -52,6 +54,35 @@ export const DEFAULT_PROVIDERS: BifrostProvider[] = [ */ export const BIFROST_VK_PROVIDER_REFRESH_MS = 24 * 60 * 60 * 1000; +/** + * Shorter re-check window used when the model a caller asked for maps + * to a provider the cached VK's grant snapshot lacks — or when the row + * has no snapshot yet. Without it, a gateway that gains a provider + * (xai via stakgraph#1673) would take a full + * `BIFROST_VK_PROVIDER_REFRESH_MS` to reach a cached VK, and every + * call for that provider in between would fall back to the direct + * key. A persistent miss costs one GET (+ one PUT) per member per + * window — bounded, and only while someone keeps asking for the + * missing provider. + */ +export const BIFROST_VK_PROVIDER_MISS_REFRESH_MS = 15 * 60 * 1000; + +/** + * aieo provider prefix (the `provider/` half of a Hive model id) → + * Bifrost provider id, i.e. the name a VK `provider_configs[].provider` + * carries. Only `google` differs: aieo speaks Gemini through the + * `google` SDK, Bifrost calls the provider `gemini`. Typed against + * aieo's union so a new aieo provider fails to compile here until it + * is mapped. + */ +export const AIEO_TO_BIFROST_PROVIDER: Record = { + anthropic: "anthropic", + google: "gemini", + openai: "openai", + openrouter: "openrouter", + xai: "xai", +}; + /** Default Bifrost admin-API port. Bifrost serves admin on 8181. */ export const DEFAULT_BIFROST_PORT = 8181; diff --git a/src/services/bifrost/orchestrator.ts b/src/services/bifrost/orchestrator.ts index 8edb0b2fdc..19a45180d1 100644 --- a/src/services/bifrost/orchestrator.ts +++ b/src/services/bifrost/orchestrator.ts @@ -37,7 +37,7 @@ import { const PUBLIC_VIEWER_USER_ID = "__public_viewer__"; export { BIFROST_AGENT_NAMES, type BifrostAgentName } from "./agent-names"; -import { BIFROST_AGENT_NAMES, type BifrostAgentName } from "./agent-names"; +import type { BifrostAgentName } from "./agent-names"; /** * Master Bifrost reconciler — the **only** function callers should @@ -150,6 +150,29 @@ export async function getBifrostForLLM( const vk = await runVKReconcile(workspaceAuth, opts.model); if (!vk) return undefined; + // 3b. Grant check. A VK only reaches the providers it was granted + // (DEFAULT_PROVIDERS ∩ what that gateway has configured — see the + // reconciler). When the snapshot says this model's provider is not + // among them (a swarm whose gateway has no XAI_API_KEY yet, say), + // the gateway would reject the call and nothing downstream retries, + // so hand the caller back to its direct provider key instead. + // Unknown grants (no snapshot yet) keep the fail-open routing. + if (vk.modelProviderGranted === false) { + logger.info( + "Bifrost VK has no grant for the model's provider; falling back to direct key", + "BIFROST_VK", + { + workspaceId: workspaceAuth.workspaceId, + userId: workspaceAuth.userId, + agentName: opts.agentName, + model: opts.model ?? null, + modelProvider: vk.modelProvider, + providers: vk.providers ?? null, + }, + ); + return undefined; + } + // 4. Mint a macaroon for the x-macaroon header (phase 4 shadow // mode). Failure here MUST NOT block the LLM call — we return the // VK shape with empty headers and the call proceeds without the @@ -215,7 +238,16 @@ async function runAgentCatalogReconcile(workspaceId: string, userId?: string): P async function runVKReconcile( workspaceAuth: Required>, model: string | undefined, -): Promise<{ apiKey: string; baseUrl: string } | undefined> { +): Promise< + | { + apiKey: string; + baseUrl: string; + modelProvider: string; + providers?: string[]; + modelProviderGranted?: boolean; + } + | undefined +> { try { const { reconcileBifrostVK } = await import("./reconciler"); const result = await reconcileBifrostVK( @@ -223,7 +255,13 @@ async function runVKReconcile( workspaceAuth.userId, { model }, ); - return { apiKey: result.vkValue, baseUrl: result.baseUrl }; + return { + apiKey: result.vkValue, + baseUrl: result.baseUrl, + modelProvider: result.modelProvider, + providers: result.providers, + modelProviderGranted: result.modelProviderGranted, + }; } catch (err) { logger.warn( "Bifrost VK reconcile failed; falling back to default LLM key", diff --git a/src/services/bifrost/reconciler.ts b/src/services/bifrost/reconciler.ts index dbbd678cac..eaeaca861e 100644 --- a/src/services/bifrost/reconciler.ts +++ b/src/services/bifrost/reconciler.ts @@ -2,14 +2,16 @@ import { db } from "@/lib/db"; import { EncryptionService } from "@/lib/encryption"; import { logger } from "@/lib/logger"; import { withLock } from "@/lib/locks/redis-lock"; -import { gatewayUrlForModel } from "aieo"; +import { gatewayUrlForModel, getProviderForModel } from "aieo"; import { BifrostClient, BifrostHttpError } from "./BifrostClient"; import { + AIEO_TO_BIFROST_PROVIDER, BIFROST_LOCK_ACQUIRE_TIMEOUT_MS, BIFROST_LOCK_PREFIX, BIFROST_LOCK_TTL_MS, BIFROST_LOG_TAG, + BIFROST_VK_PROVIDER_MISS_REFRESH_MS, BIFROST_VK_PROVIDER_REFRESH_MS, DEFAULT_BUDGET_RESET_DURATION, DEFAULT_CUSTOMER_BUDGET_USD, @@ -47,7 +49,9 @@ import type { * cached VK on `WorkspaceMember` without talking to Bifrost — except * once per `BIFROST_VK_PROVIDER_REFRESH_MS`, when the VK's provider * grants are re-checked against the gateway (see "Provider grants" - * below). + * below). The grants last observed are snapshotted on the row so a + * call for a provider the VK doesn't carry can be steered back to + * the caller's direct key instead of failing at the gateway. * * See `gateway/plans/phase-1-reconciler.md`. */ @@ -114,6 +118,7 @@ async function doReconcile( bifrostVkId: true, bifrostCustomerId: true, bifrostSyncedAt: true, + bifrostVkProviders: true, user: { select: { githubAuth: { select: { githubUsername: true } }, @@ -140,6 +145,7 @@ async function doReconcile( // downstream agent will call directly, so it needs the provider // suffix already applied. const llmBaseUrl = gatewayUrlForModel(options.model, baseCreds.baseUrl); + const modelProvider = bifrostProviderForModel(options.model); // Constructing the client is free (no I/O). It's only exercised on // the Bifrost paths below — including the cached path's periodic @@ -170,13 +176,37 @@ async function doReconcile( ); } if (vkValue !== undefined) { + // Grant snapshot from the last create / refresh. Empty means + // unknown (row predates the column, or the gateway couldn't be + // read) — never "no grants". + let providers = knownGrants(member.bifrostVkProviders); // The cache never talks to Bifrost, so a provider added to // DEFAULT_PROVIDERS (or newly configured on this gateway) would // otherwise never reach a VK minted before it. Re-check the - // grants once per refresh window. Best-effort — never blocks - // or fails the call. - if (isProviderGrantStale(member.bifrostSyncedAt)) { - await refreshProviderGrants(client, member.bifrostVkId, member.id); + // grants once per refresh window — and sooner, once per miss + // window, when the caller's model needs a provider the snapshot + // lacks (or there is no snapshot), so a gateway that just gained + // a provider flips its cached VKs in minutes rather than a day. + // Best-effort — never blocks or fails the call. + const missing = + providers === undefined || !providers.includes(modelProvider); + if ( + isProviderGrantStale( + member.bifrostSyncedAt, + BIFROST_VK_PROVIDER_REFRESH_MS, + ) || + (missing && + isProviderGrantStale( + member.bifrostSyncedAt, + BIFROST_VK_PROVIDER_MISS_REFRESH_MS, + )) + ) { + const refreshed = await refreshProviderGrants( + client, + member.bifrostVkId, + member.id, + ); + if (refreshed) providers = refreshed; } return { workspaceId, @@ -186,6 +216,9 @@ async function doReconcile( vkValue, baseUrl: llmBaseUrl, created: false, + modelProvider, + providers, + modelProviderGranted: isGranted(providers, modelProvider), }; } } @@ -200,7 +233,7 @@ async function doReconcile( if (createdCustomer) created = true; const providers = await desiredProviders(client); - const { virtualKey, createdVk } = await ensureVirtualKey( + const { virtualKey, createdVk, granted } = await ensureVirtualKey( client, bifrostName, customer.id, @@ -219,6 +252,9 @@ async function doReconcile( bifrostVkId: virtualKey.id, bifrostCustomerId: customer.id, bifrostSyncedAt: new Date(), + // Empty when the grants couldn't be observed — read as "unknown" + // by the cached path, which then refreshes on its next miss. + bifrostVkProviders: granted ?? [], }, }); @@ -239,6 +275,9 @@ async function doReconcile( vkValue: virtualKey.value, baseUrl: llmBaseUrl, created, + modelProvider, + providers: granted, + modelProviderGranted: isGranted(granted, modelProvider), }; } @@ -310,14 +349,21 @@ async function ensureVirtualKey( name: string, customerId: string, providers: BifrostProvider[], -): Promise<{ virtualKey: BifrostVirtualKey; createdVk: boolean }> { +): Promise<{ + virtualKey: BifrostVirtualKey; + createdVk: boolean; + /** Providers the VK carries after this call; undefined when unobserved. */ + granted?: string[]; +}> { const existing = await findExactVirtualKey(client, name, customerId); if (existing) { + let granted = grantsOf(existing); // Non-fatal: the VK already works for the providers it has, and // failing here would make the orchestrator drop this call to the // swarm's default key over a grant it may not even need. try { - await topUpProviderGrants(client, existing, providers); + const topped = await topUpProviderGrants(client, existing, providers); + if (topped.granted) granted = topped.granted; } catch (err) { logger.warn( "Bifrost VK provider top-up failed; using VK as is", @@ -330,7 +376,7 @@ async function ensureVirtualKey( }, ); } - return { virtualKey: existing, createdVk: false }; + return { virtualKey: existing, createdVk: false, granted }; } try { @@ -354,7 +400,9 @@ async function ensureVirtualKey( key_ids: ["*"], })), }); - return { virtualKey: created.virtual_key, createdVk: true }; + // What we asked for is what the VK carries: Bifrost 400s the whole + // create on a provider it lacks rather than dropping that entry. + return { virtualKey: created.virtual_key, createdVk: true, granted: providers }; } catch (err) { // VK names are uniquely indexed at the DB level (plan §4). On a // dup-key race, read back per the plan. @@ -366,7 +414,11 @@ async function ensureVirtualKey( BIFROST_LOG_TAG, { name, customerId, picked: readback.id }, ); - return { virtualKey: readback, createdVk: false }; + return { + virtualKey: readback, + createdVk: false, + granted: grantsOf(readback), + }; } } throw err; @@ -447,7 +499,8 @@ async function desiredProviders( /** * Grant the VK any of `desired` it doesn't already carry. Returns the - * providers added (empty when nothing changed). + * providers added (empty when nothing changed) and the providers the + * VK carries afterwards (undefined when the response didn't say). * * Bifrost's PUT replaces the provider_configs set wholesale: entries * with an `id` update in place, entries without one are created, and @@ -459,21 +512,23 @@ async function topUpProviderGrants( client: BifrostClient, vk: BifrostVirtualKey, desired: BifrostProvider[], -): Promise { +): Promise<{ added: BifrostProvider[]; granted?: string[] }> { const existing = vk.provider_configs; // Without the hydrated set (ids included) we can't re-send it, and - // a PUT that omits a config deletes it. Leave the VK alone. + // a PUT that omits a config deletes it. Leave the VK alone — but the + // provider names are still a faithful reading of what it carries. if (!Array.isArray(existing) || existing.some((pc) => pc.id == null)) { logger.warn( "Bifrost VK response lacks hydrated provider_configs; skipping grant top-up", BIFROST_LOG_TAG, { vkId: vk.id }, ); - return []; + return { added: [], granted: grantsOf(vk) }; } - const have = new Set(existing.map((pc) => pc.provider)); - const missing = desired.filter((p) => !have.has(p)); - if (missing.length === 0) return []; + const have = existing.map((pc) => pc.provider); + const haveSet = new Set(have); + const missing = desired.filter((p) => !haveSet.has(p)); + if (missing.length === 0) return { added: [], granted: have }; await client.updateVirtualKey(vk.id, { provider_configs: [ @@ -489,7 +544,7 @@ async function topUpProviderGrants( vkId: vk.id, added: missing, }); - return missing; + return { added: missing, granted: [...have, ...missing] }; } /** @@ -516,7 +571,9 @@ function carryProviderConfig(pc: BifrostProviderConfig) { /** * Cached-path companion to `topUpProviderGrants`: fetch the cached VK * from the gateway, grant anything it's missing, then stamp - * `bifrostSyncedAt` so the next check is a refresh window away. + * `bifrostSyncedAt` (and the grant snapshot, when observed) so the + * next check is a refresh window away. Returns the grants observed, + * or undefined when the gateway couldn't be read. * * Best-effort throughout. Any failure is logged and the cached VK is * served as is — an LLM call must never fail because a grant refresh @@ -527,11 +584,12 @@ async function refreshProviderGrants( client: BifrostClient, vkId: string, memberId: string, -): Promise { +): Promise { + let granted: string[] | undefined; try { const desired = await desiredProviders(client); const { virtual_key } = await client.getVirtualKey(vkId); - await topUpProviderGrants(client, virtual_key, desired); + granted = (await topUpProviderGrants(client, virtual_key, desired)).granted; } catch (err) { logger.warn( "Bifrost VK grant refresh failed; serving cached VK", @@ -542,7 +600,12 @@ async function refreshProviderGrants( try { await db.workspaceMember.update({ where: { id: memberId }, - data: { bifrostSyncedAt: new Date() }, + data: { + bifrostSyncedAt: new Date(), + // Only overwrite the snapshot when this refresh actually saw + // the VK; a failed read must not blank a good one. + ...(granted ? { bifrostVkProviders: granted } : {}), + }, }); } catch (err) { logger.warn( @@ -551,11 +614,49 @@ async function refreshProviderGrants( { memberId, error: err instanceof Error ? err.message : String(err) }, ); } + return granted; } -function isProviderGrantStale(syncedAt: Date | null | undefined): boolean { +function isProviderGrantStale( + syncedAt: Date | null | undefined, + windowMs: number, +): boolean { if (!syncedAt) return true; - return Date.now() - syncedAt.getTime() > BIFROST_VK_PROVIDER_REFRESH_MS; + return Date.now() - syncedAt.getTime() > windowMs; +} + +/** + * Bifrost provider id for the model a caller intends to use — the + * same prefix resolution `gatewayUrlForModel` applies to the base URL + * (shortcuts like `"grok"`, namespaced ids like `"xai/grok-4.6"`, bare + * ids, anthropic when omitted), mapped through + * `AIEO_TO_BIFROST_PROVIDER` (`google` → `gemini`). + */ +export function bifrostProviderForModel( + model: string | undefined, +): BifrostProvider { + return AIEO_TO_BIFROST_PROVIDER[getProviderForModel(model)]; +} + +/** Snapshot → known grants. The empty default means "unknown", never "none". */ +function knownGrants( + snapshot: string[] | null | undefined, +): string[] | undefined { + return snapshot && snapshot.length > 0 ? snapshot : undefined; +} + +/** Provider names a VK response carries, or undefined when it doesn't say. */ +function grantsOf(vk: BifrostVirtualKey): string[] | undefined { + return Array.isArray(vk.provider_configs) + ? vk.provider_configs.map((pc) => pc.provider) + : undefined; +} + +function isGranted( + providers: string[] | undefined, + provider: string, +): boolean | undefined { + return providers === undefined ? undefined : providers.includes(provider); } /** diff --git a/src/services/bifrost/types.ts b/src/services/bifrost/types.ts index fe52aa113d..a31bb041df 100644 --- a/src/services/bifrost/types.ts +++ b/src/services/bifrost/types.ts @@ -167,6 +167,28 @@ export interface ReconcileResult { baseUrl: string; /** True iff Customer or VK was created during this call (audit signal). */ created: boolean; + /** + * Bifrost provider the requested model maps to (`options.model`, + * anthropic when omitted) — the same prefix resolution that picked + * the `baseUrl` suffix, mapped through `AIEO_TO_BIFROST_PROVIDER` + * (`gemini` for `google/*` models). + */ + modelProvider: BifrostProvider; + /** + * The VK's provider grants as last observed: the create response, or + * the snapshot on `WorkspaceMember.bifrostVkProviders` (refreshed + * once per window, sooner on a miss). `undefined` when unknown — the + * row has no snapshot yet and the gateway couldn't be read this call. + */ + providers?: string[]; + /** + * `false` when `providers` is known and lacks `modelProvider`. The + * orchestrator then returns `undefined` so the caller uses its + * direct provider key instead of sending a call the gateway would + * reject. `true` when granted; `undefined` when grants are unknown + * (fail open — route through Bifrost as before). + */ + modelProviderGranted?: boolean; } export interface BifrostAdminCreds { diff --git a/src/services/task-workflow.ts b/src/services/task-workflow.ts index d328c48f28..3b0f2c7886 100644 --- a/src/services/task-workflow.ts +++ b/src/services/task-workflow.ts @@ -872,31 +872,23 @@ export async function callStakworkAPI(params: { // tuning of ttlSeconds / maxCostUsd / maxSteps is intentionally // deferred to a follow-up so this initial wiring stays small. // - // xAI bypass: `DEFAULT_PROVIDERS` now lists "xai" and the VK - // reconciler grants it on any gateway that has it configured, but - // the swarm gateways don't yet carry an XAI_API_KEY (sphinx-swarm's - // bifrost config is the missing piece), so routing an `xai/*` - // selection through Bifrost today would still fail on most swarms. - // Skip the Bifrost call entirely for xai/* and fall through to the - // direct `vars.apiKey` (XAI_API_KEY) resolved above. Remove this - // bypass once the swarm gateways ship with an xai key (aieo already - // maps xai onto the /openai/v1 gateway path) — 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, - }, - ); + // Every provider rides Bifrost here, xai included — swarm gateways + // carry an XAI_API_KEY since sphinx-swarm's "xai api key in bifrost". + // A swarm whose gateway still lacks the model's provider is handled + // inside the orchestrator: it returns `undefined` when the VK has no + // grant for that provider, so `vars.apiKey` stays the direct key + // resolved above for that call. + 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, + }, + ); if (bifrost) { vars.apiKey = bifrost.apiKey; vars.baseUrl = bifrost.baseUrl;