Skip to content

Commit 959eee4

Browse files
authored
feat(provider): route the managed wallet through OpenRouter only (#110)
When the LLM spend toggle is explicitly "managed", wallet inference now flows through OpenRouter exclusively — the first-party managed proxies (anthropic / openai / google), which openscience had fanned the wallet across via managedProxyKey, are taken out of the managed path. OpenRouter is the one gateway with a single unified reasoning stream, so this both implements "wallet ⇒ OpenRouter, no other providers" and sets up the reasoning-trace fix (follow-up PR). Two pure helpers gate everything: - managedRoutesOpenRouterOnly(config): true only for billing.llm==='managed'. - managedProviderAllowed(id): openrouter + the hosted synsci demo. Wired into (1) managedProxyKey — never attach the thk_ wallet token to a first-party proxy; and (2) isProviderAllowed — drop non-OpenRouter/non-synsci providers from a managed session, which transitively makes defaultModel() and getSmallModel() OpenRouter-only (they read the filtered state). BYOK and the legacy auto-detect path (billing.llm unset/null/'byok') are untouched — the existing managed-proxy-forwarding tests, which never set the toggle, still pass. pinByokToPublicEndpoint and the managed-key sanity check are unchanged.
1 parent dececb9 commit 959eee4

2 files changed

Lines changed: 189 additions & 5 deletions

File tree

backend/cli/src/provider/provider.ts

Lines changed: 40 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -93,10 +93,16 @@ export namespace Provider {
9393
// managed token and 401 the proxy ("thk_* token not found"). Returns {} unless
9494
// both managed spend is on AND the baseURL is the Atlas proxy — genuine BYOK
9595
// (no proxy URL) is untouched and hits the provider directly.
96-
async function managedProxyKey(baseURL: unknown): Promise<{ apiKey?: string }> {
96+
async function managedProxyKey(providerID: string, baseURL: unknown): Promise<{ apiKey?: string }> {
9797
const managed =
9898
isAtlasProxyBaseURL(baseURL) && (await Config.get().catch(() => undefined))?.billing?.llm === "managed"
9999
if (!managed) return {}
100+
// Managed wallet routes through OpenRouter ONLY: never attach the wallet's
101+
// thk_ token to a first-party proxy (anthropic / openai / google). Those
102+
// providers are also dropped from availability (see isProviderAllowed), so
103+
// this is belt-and-suspenders — a managed token can only ever reach the
104+
// sanctioned OpenRouter (or hosted synsci) route.
105+
if (!managedProviderAllowed(providerID)) return {}
100106
const session = await OpenScience.getSession().catch(() => null)
101107
return session?.api_key ? { apiKey: session.api_key } : {}
102108
}
@@ -124,6 +130,26 @@ export namespace Provider {
124130
return typeof key === "string" && key.length > 0 && key !== "public" && !isAtlasApiKey(key)
125131
}
126132

133+
/** Managed wallet ⇒ OpenRouter-only routing.
134+
*
135+
* When the LLM spend toggle is explicitly "managed", every wallet inference
136+
* call flows through OpenRouter — the one gateway whose stream exposes a
137+
* single, unified reasoning format (`reasoning` / `reasoning_details`). The
138+
* first-party managed proxies (anthropic / openai / google), each with a
139+
* different reasoning shape, are taken out of the managed path entirely; the
140+
* hosted zero-cost `synsci` demo provider is kept. BYOK and the legacy
141+
* auto-detect path (`billing.llm` unset / null / "byok") are UNTOUCHED —
142+
* this only fires on an explicit managed-wallet opt-in. Pure + sync. */
143+
export function managedRoutesOpenRouterOnly(config: Config.Info): boolean {
144+
return config.billing?.llm === "managed"
145+
}
146+
147+
/** Providers a managed (OpenRouter-only) wallet session may load: OpenRouter
148+
* for all real inference, plus the hosted `synsci` demo. Pure + sync. */
149+
export function managedProviderAllowed(providerID: string): boolean {
150+
return providerID === "openrouter" || providerID.startsWith("synsci")
151+
}
152+
127153
/**
128154
* Blocker (c) — the inverse of requireAtlasProxyForManagedKey.
129155
*
@@ -162,7 +188,7 @@ export namespace Provider {
162188
return {
163189
autoload: false,
164190
options: {
165-
...(baseURL ? { baseURL, ...(await managedProxyKey(baseURL)) } : {}),
191+
...(baseURL ? { baseURL, ...(await managedProxyKey("anthropic", baseURL)) } : {}),
166192
headers: {
167193
"anthropic-beta":
168194
"claude-code-20250219,interleaved-thinking-2025-05-14,fine-grained-tool-streaming-2025-05-14",
@@ -199,7 +225,7 @@ export namespace Provider {
199225
async getModel(sdk: any, modelID: string, _options?: Record<string, any>) {
200226
return sdk.responses(modelID)
201227
},
202-
options: baseURL ? { baseURL, ...(await managedProxyKey(baseURL)) } : {},
228+
options: baseURL ? { baseURL, ...(await managedProxyKey("openai", baseURL)) } : {},
203229
}
204230
},
205231
"github-copilot": async () => {
@@ -405,7 +431,7 @@ export namespace Provider {
405431
return {
406432
autoload: false,
407433
options: {
408-
...(baseURL ? { baseURL, ...(await managedProxyKey(baseURL)) } : {}),
434+
...(baseURL ? { baseURL, ...(await managedProxyKey("openrouter", baseURL)) } : {}),
409435
headers: {
410436
"HTTP-Referer": "https://syntheticsciences.ai/",
411437
"X-Title": "synsci",
@@ -598,7 +624,7 @@ export namespace Provider {
598624
const baseURL = Env.get("GOOGLE_GENERATIVE_AI_BASE_URL") ?? Env.get("GEMINI_BASE_URL")
599625
return {
600626
autoload: false,
601-
options: baseURL ? { baseURL, ...(await managedProxyKey(baseURL)) } : {},
627+
options: baseURL ? { baseURL, ...(await managedProxyKey("google", baseURL)) } : {},
602628
}
603629
},
604630
}
@@ -838,8 +864,17 @@ export namespace Provider {
838864

839865
const disabled = new Set(config.disabled_providers ?? [])
840866
const enabled = config.enabled_providers ? new Set(config.enabled_providers) : null
867+
// Managed wallet ⇒ OpenRouter-only. Drop every other provider (the
868+
// first-party managed proxies included) from a managed session so wallet
869+
// inference can only flow through OpenRouter's unified reasoning stream,
870+
// plus the hosted zero-cost demo. Gated on the explicit toggle, so BYOK and
871+
// legacy auto-detect sessions see every provider exactly as before. This is
872+
// the single seam that makes defaultModel()/getSmallModel() managed-safe:
873+
// both read the filtered state, so they can only resolve openrouter/synsci.
874+
const managedOpenRouterOnly = managedRoutesOpenRouterOnly(config)
841875

842876
function isProviderAllowed(providerID: string): boolean {
877+
if (managedOpenRouterOnly && !managedProviderAllowed(providerID)) return false
843878
if (enabled && !enabled.has(providerID)) return false
844879
if (disabled.has(providerID)) return false
845880
return true
Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
import { test, expect, describe, mock } from "bun:test"
2+
3+
// Mock BunProc + default auth plugins so importing Provider never shells out or
4+
// hits the network (mirrors test/provider/provider.test.ts).
5+
mock.module("../../src/bun/index", () => ({
6+
BunProc: {
7+
install: async (pkg: string, _version?: string) => {
8+
const at = pkg.lastIndexOf("@")
9+
return at > 0 ? pkg.substring(0, at) : pkg
10+
},
11+
run: async () => {
12+
throw new Error("BunProc.run should not be called in tests")
13+
},
14+
which: () => process.execPath,
15+
InstallFailedError: class extends Error {},
16+
},
17+
}))
18+
const mockPlugin = () => ({})
19+
mock.module("openscience-copilot-auth", () => ({ default: mockPlugin }))
20+
mock.module("openscience-anthropic-auth", () => ({ default: mockPlugin }))
21+
mock.module("@gitlab/openscience-gitlab-auth", () => ({ default: mockPlugin }))
22+
23+
import { tmpdir } from "../fixture/fixture"
24+
import { Instance } from "../../src/project/instance"
25+
import { Provider } from "../../src/provider/provider"
26+
import { Env } from "../../src/env"
27+
28+
function clearManagedLLMEnv() {
29+
for (const key of [
30+
"ANTHROPIC_API_KEY",
31+
"ANTHROPIC_BASE_URL",
32+
"OPENAI_API_KEY",
33+
"OPENAI_BASE_URL",
34+
"GOOGLE_GENERATIVE_AI_API_KEY",
35+
"GOOGLE_GENERATIVE_AI_BASE_URL",
36+
"GEMINI_API_KEY",
37+
"GEMINI_BASE_URL",
38+
"OPENROUTER_API_KEY",
39+
"OPENROUTER_BASE_URL",
40+
]) {
41+
Env.remove(key)
42+
}
43+
}
44+
45+
const PROXY = "https://atlas.test/api/llm/proxy"
46+
47+
// ── Pure decision helpers ────────────────────────────────────────────────────
48+
49+
describe("Provider.managedRoutesOpenRouterOnly (pure)", () => {
50+
test("true only when billing.llm === 'managed'", () => {
51+
expect(Provider.managedRoutesOpenRouterOnly({ billing: { llm: "managed" } } as any)).toBe(true)
52+
})
53+
test("false for explicit byok", () => {
54+
expect(Provider.managedRoutesOpenRouterOnly({ billing: { llm: "byok" } } as any)).toBe(false)
55+
})
56+
test("false for auto-detect (unset / empty billing / null)", () => {
57+
expect(Provider.managedRoutesOpenRouterOnly({} as any)).toBe(false)
58+
expect(Provider.managedRoutesOpenRouterOnly({ billing: {} } as any)).toBe(false)
59+
expect(Provider.managedRoutesOpenRouterOnly({ billing: { llm: null } } as any)).toBe(false)
60+
})
61+
})
62+
63+
describe("Provider.managedProviderAllowed (pure)", () => {
64+
test("OpenRouter and the hosted synsci demo are the only allowed providers", () => {
65+
expect(Provider.managedProviderAllowed("openrouter")).toBe(true)
66+
expect(Provider.managedProviderAllowed("synsci")).toBe(true)
67+
expect(Provider.managedProviderAllowed("synsci-hosted")).toBe(true)
68+
})
69+
test("first-party managed proxies + everything else are rejected", () => {
70+
for (const id of ["anthropic", "openai", "google", "openai-codex", "github-copilot", "gateway", "azure"]) {
71+
expect(Provider.managedProviderAllowed(id)).toBe(false)
72+
}
73+
})
74+
})
75+
76+
// ── Availability filter (hermetic, catalog-backed) ───────────────────────────
77+
78+
describe("managed session availability", () => {
79+
test("managed ⇒ only OpenRouter (+ demo) load; first-party proxies are dropped", async () => {
80+
await using tmp = await tmpdir({ config: { billing: { llm: "managed" } } })
81+
await Instance.provide({
82+
directory: tmp.path,
83+
init: async () => {
84+
clearManagedLLMEnv()
85+
Env.set("ANTHROPIC_API_KEY", "thk_anthropic")
86+
Env.set("ANTHROPIC_BASE_URL", `${PROXY}/anthropic/v1`)
87+
Env.set("OPENAI_API_KEY", "thk_openai")
88+
Env.set("OPENAI_BASE_URL", `${PROXY}/openai/v1`)
89+
Env.set("GOOGLE_GENERATIVE_AI_API_KEY", "thk_google")
90+
Env.set("GOOGLE_GENERATIVE_AI_BASE_URL", `${PROXY}/gemini/v1beta`)
91+
Env.set("OPENROUTER_API_KEY", "thk_openrouter")
92+
Env.set("OPENROUTER_BASE_URL", `${PROXY}/openrouter/v1`)
93+
Provider.invalidate()
94+
},
95+
fn: async () => {
96+
const providers = await Provider.list()
97+
expect(providers["openrouter"]).toBeDefined()
98+
expect(providers["openrouter"].options.baseURL).toBe(`${PROXY}/openrouter/v1`)
99+
expect(providers["anthropic"]).toBeUndefined()
100+
expect(providers["openai"]).toBeUndefined()
101+
expect(providers["google"]).toBeUndefined()
102+
},
103+
})
104+
})
105+
106+
test("BYOK (managed off): anthropic keeps its public endpoint and no wallet token", async () => {
107+
await using tmp = await tmpdir({ config: { billing: { llm: "byok" } } })
108+
await Instance.provide({
109+
directory: tmp.path,
110+
init: async () => {
111+
clearManagedLLMEnv()
112+
Env.set("ANTHROPIC_API_KEY", "sk-ant-byok-key")
113+
Provider.invalidate()
114+
},
115+
fn: async () => {
116+
const providers = await Provider.list()
117+
const anthropic = providers["anthropic"]
118+
expect(anthropic).toBeDefined()
119+
// No Atlas proxy baseURL injected → getSDK falls back to the public
120+
// model.api.url; the BYOK key is the only credential.
121+
expect(anthropic.options.baseURL).toBeUndefined()
122+
expect(anthropic.options.apiKey).toBeUndefined()
123+
expect(anthropic.key).toBe("sk-ant-byok-key")
124+
},
125+
})
126+
})
127+
128+
test("legacy auto-detect (thk_ present, billing.llm unset) is unchanged — proxies still load", async () => {
129+
await using tmp = await tmpdir({ config: {} })
130+
await Instance.provide({
131+
directory: tmp.path,
132+
init: async () => {
133+
clearManagedLLMEnv()
134+
Env.set("ANTHROPIC_API_KEY", "thk_anthropic")
135+
Env.set("ANTHROPIC_BASE_URL", `${PROXY}/anthropic/v1`)
136+
Env.set("OPENROUTER_API_KEY", "thk_openrouter")
137+
Env.set("OPENROUTER_BASE_URL", `${PROXY}/openrouter/v1`)
138+
Provider.invalidate()
139+
},
140+
fn: async () => {
141+
const providers = await Provider.list()
142+
// Gating is scoped to the explicit toggle: without it, nothing is dropped.
143+
expect(providers["anthropic"]).toBeDefined()
144+
expect(providers["anthropic"].options.baseURL).toBe(`${PROXY}/anthropic/v1`)
145+
expect(providers["openrouter"]).toBeDefined()
146+
},
147+
})
148+
})
149+
})

0 commit comments

Comments
 (0)