1
0

provider-google-vertex.test.ts 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380
  1. import { AISDK } from "@opencode-ai/core/aisdk"
  2. import { describe, expect, mock } from "bun:test"
  3. import { Effect } from "effect"
  4. import { Catalog } from "@opencode-ai/core/catalog"
  5. import { Model } from "@opencode-ai/core/model"
  6. import { Plugin } from "@opencode-ai/core/plugin"
  7. import { PluginHost } from "@opencode-ai/core/plugin/host"
  8. import { GoogleVertexPlugin } from "@opencode-ai/core/plugin/provider/google-vertex"
  9. import { Provider } from "@opencode-ai/core/provider"
  10. import type { LanguageModelV3 } from "@ai-sdk/provider"
  11. import { testEffect } from "../lib/effect"
  12. import { PluginTestLayer } from "./fixture"
  13. const vertexOptions: Record<string, any>[] = []
  14. const googleAuthOptions: Record<string, any>[] = []
  15. const it = testEffect(PluginTestLayer)
  16. const addPlugin = Effect.fn(function* () {
  17. const plugin = yield* Plugin.Service
  18. const aisdk = yield* AISDK.Service
  19. const host = yield* PluginHost.make(plugin)
  20. yield* GoogleVertexPlugin.effect(host)
  21. })
  22. function required<T>(value: T | undefined): T {
  23. if (value === undefined) throw new Error("Expected value")
  24. return value
  25. }
  26. function withEnv<A, E, R>(vars: Record<string, string | undefined>, effect: () => Effect.Effect<A, E, R>) {
  27. return Effect.acquireUseRelease(
  28. Effect.sync(() => {
  29. const previous = Object.fromEntries(Object.keys(vars).map((key) => [key, process.env[key]]))
  30. Object.entries(vars).forEach(([key, value]) => {
  31. if (value === undefined) delete process.env[key]
  32. else process.env[key] = value
  33. })
  34. return previous
  35. }),
  36. effect,
  37. (previous) =>
  38. Effect.sync(() =>
  39. Object.entries(previous).forEach(([key, value]) => {
  40. if (value === undefined) delete process.env[key]
  41. else process.env[key] = value
  42. }),
  43. ),
  44. )
  45. }
  46. function fakeSelectorSdk(calls: string[]) {
  47. const make = (method: string) => (id: string) => {
  48. calls.push(`${method}:${id}`)
  49. return { modelId: id, provider: method, specificationVersion: "v3" } as unknown as LanguageModelV3
  50. }
  51. return {
  52. responses: make("responses"),
  53. messages: make("messages"),
  54. chat: make("chat"),
  55. languageModel: make("languageModel"),
  56. }
  57. }
  58. void mock.module("@ai-sdk/google-vertex", () => ({
  59. createVertex: (options: Record<string, any>) => {
  60. vertexOptions.push(options)
  61. return {
  62. languageModel: (modelID: string) => ({ modelID, provider: "google-vertex", specificationVersion: "v3" }),
  63. }
  64. },
  65. }))
  66. void mock.module("google-auth-library", () => ({
  67. GoogleAuth: class {
  68. constructor(options: Record<string, any>) {
  69. googleAuthOptions.push(options)
  70. }
  71. async getClient() {
  72. return {
  73. async getAccessToken() {
  74. return { token: "vertex-token" }
  75. },
  76. }
  77. }
  78. },
  79. }))
  80. describe("GoogleVertexPlugin", () => {
  81. it.effect("ignores OpenAI-compatible providers that are not Google Vertex", () =>
  82. Effect.gen(function* () {
  83. const catalog = yield* Catalog.Service
  84. yield* catalog.transform((catalog) =>
  85. catalog.provider.update(Provider.ID.opencode, (provider) => {
  86. provider.package = Provider.aisdk("@ai-sdk/openai-compatible")
  87. provider.settings = { ...provider.settings, baseURL: "https://opencode.ai/zen/v1" }
  88. }),
  89. )
  90. yield* addPlugin()
  91. const provider = required(yield* catalog.provider.get(Provider.ID.opencode))
  92. expect(provider.settings).toEqual({ baseURL: "https://opencode.ai/zen/v1" })
  93. }),
  94. )
  95. it.effect("resolves project and location from env using legacy precedence", () =>
  96. withEnv(
  97. {
  98. GOOGLE_CLOUD_PROJECT: "google-cloud-project",
  99. GCP_PROJECT: "gcp-project",
  100. GCLOUD_PROJECT: "gcloud-project",
  101. GOOGLE_VERTEX_LOCATION: "google-vertex-location",
  102. GOOGLE_CLOUD_LOCATION: "google-cloud-location",
  103. VERTEX_LOCATION: "vertex-location",
  104. },
  105. () =>
  106. Effect.gen(function* () {
  107. const catalog = yield* Catalog.Service
  108. yield* catalog.transform((catalog) =>
  109. catalog.provider.update(Provider.ID.make("google-vertex"), (provider) => {
  110. provider.package = Provider.aisdk("@ai-sdk/openai-compatible")
  111. provider.settings = {
  112. ...provider.settings,
  113. baseURL:
  114. "https://${GOOGLE_VERTEX_ENDPOINT}/v1/projects/${GOOGLE_VERTEX_PROJECT}/locations/${GOOGLE_VERTEX_LOCATION}",
  115. }
  116. }),
  117. )
  118. yield* addPlugin()
  119. const provider = required(yield* catalog.provider.get(Provider.ID.make("google-vertex")))
  120. expect(provider.settings?.project).toBe("google-cloud-project")
  121. expect(provider.settings?.location).toBe("google-vertex-location")
  122. expect(provider).toMatchObject({
  123. package: "aisdk:@ai-sdk/openai-compatible",
  124. settings: {
  125. baseURL:
  126. "https://google-vertex-location-aiplatform.googleapis.com/v1/projects/google-cloud-project/locations/google-vertex-location",
  127. },
  128. })
  129. }),
  130. ),
  131. )
  132. it.effect("resolves the advertised GOOGLE_VERTEX_PROJECT env for provider updates and SDKs", () =>
  133. withEnv(
  134. {
  135. GOOGLE_VERTEX_PROJECT: "vertex-project",
  136. GOOGLE_CLOUD_PROJECT: undefined,
  137. GCP_PROJECT: undefined,
  138. GCLOUD_PROJECT: undefined,
  139. GOOGLE_VERTEX_LOCATION: "europe-west4",
  140. GOOGLE_CLOUD_LOCATION: undefined,
  141. VERTEX_LOCATION: undefined,
  142. },
  143. () =>
  144. Effect.gen(function* () {
  145. vertexOptions.length = 0
  146. const plugin = yield* Plugin.Service
  147. const aisdk = yield* AISDK.Service
  148. const catalog = yield* Catalog.Service
  149. yield* catalog.transform((catalog) =>
  150. catalog.provider.update(Provider.ID.make("google-vertex"), (provider) => {
  151. provider.package = Provider.aisdk("@ai-sdk/openai-compatible")
  152. provider.settings = {
  153. ...provider.settings,
  154. baseURL:
  155. "https://${GOOGLE_VERTEX_ENDPOINT}/v1/projects/${GOOGLE_VERTEX_PROJECT}/locations/${GOOGLE_VERTEX_LOCATION}",
  156. }
  157. }),
  158. )
  159. yield* addPlugin()
  160. const provider = required(yield* catalog.provider.get(Provider.ID.make("google-vertex")))
  161. yield* aisdk.runSDK({
  162. model: Model.Info.make({
  163. ...Model.Info.default(Provider.ID.make("google-vertex"), Model.ID.make("gemini")),
  164. modelID: Model.ID.make("gemini"),
  165. package: "aisdk:@ai-sdk/google-vertex",
  166. }),
  167. package: "@ai-sdk/google-vertex",
  168. options: { name: "google-vertex" },
  169. })
  170. expect(provider.settings?.project).toBe("vertex-project")
  171. expect(provider).toMatchObject({
  172. package: "aisdk:@ai-sdk/openai-compatible",
  173. settings: {
  174. baseURL:
  175. "https://europe-west4-aiplatform.googleapis.com/v1/projects/vertex-project/locations/europe-west4",
  176. },
  177. })
  178. expect(vertexOptions[0].project).toBe("vertex-project")
  179. expect(vertexOptions[0].location).toBe("europe-west4")
  180. }),
  181. ),
  182. )
  183. it.effect("keeps configured project and location over env and uses global endpoint", () =>
  184. withEnv(
  185. {
  186. GOOGLE_CLOUD_PROJECT: "env-project",
  187. GCP_PROJECT: "env-gcp-project",
  188. GCLOUD_PROJECT: "env-gcloud-project",
  189. GOOGLE_VERTEX_LOCATION: "env-location",
  190. GOOGLE_CLOUD_LOCATION: "env-google-cloud-location",
  191. VERTEX_LOCATION: "env-vertex-location",
  192. },
  193. () =>
  194. Effect.gen(function* () {
  195. const catalog = yield* Catalog.Service
  196. yield* catalog.transform((catalog) =>
  197. catalog.provider.update(Provider.ID.make("google-vertex"), (provider) => {
  198. provider.package = Provider.aisdk("@ai-sdk/openai-compatible")
  199. provider.settings = {
  200. ...provider.settings,
  201. baseURL:
  202. "https://${GOOGLE_VERTEX_ENDPOINT}/v1/projects/${GOOGLE_VERTEX_PROJECT}/locations/${GOOGLE_VERTEX_LOCATION}",
  203. }
  204. provider.settings = { ...provider.settings, project: "config-project", location: "global" }
  205. }),
  206. )
  207. yield* addPlugin()
  208. const provider = required(yield* catalog.provider.get(Provider.ID.make("google-vertex")))
  209. expect(provider.settings?.project).toBe("config-project")
  210. expect(provider.settings?.location).toBe("global")
  211. expect(provider).toMatchObject({
  212. package: "aisdk:@ai-sdk/openai-compatible",
  213. settings: { baseURL: "https://aiplatform.googleapis.com/v1/projects/config-project/locations/global" },
  214. })
  215. }),
  216. ),
  217. )
  218. it.effect("keeps OpenAI-compatible Vertex endpoint templates regional for eu", () =>
  219. Effect.gen(function* () {
  220. const catalog = yield* Catalog.Service
  221. yield* catalog.transform((catalog) =>
  222. catalog.provider.update(Provider.ID.make("google-vertex"), (provider) => {
  223. provider.package = Provider.aisdk("@ai-sdk/openai-compatible")
  224. provider.settings = {
  225. ...provider.settings,
  226. baseURL:
  227. "https://${GOOGLE_VERTEX_ENDPOINT}/v1/projects/${GOOGLE_VERTEX_PROJECT}/locations/${GOOGLE_VERTEX_LOCATION}",
  228. }
  229. provider.settings = { ...provider.settings, project: "config-project", location: "eu" }
  230. }),
  231. )
  232. yield* addPlugin()
  233. const provider = required(yield* catalog.provider.get(Provider.ID.make("google-vertex")))
  234. expect(provider).toMatchObject({
  235. package: "aisdk:@ai-sdk/openai-compatible",
  236. settings: { baseURL: "https://eu-aiplatform.googleapis.com/v1/projects/config-project/locations/eu" },
  237. })
  238. }),
  239. )
  240. it.effect("defaults location to us-central1 when only project is configured", () =>
  241. withEnv(
  242. {
  243. GOOGLE_CLOUD_PROJECT: undefined,
  244. GCP_PROJECT: undefined,
  245. GCLOUD_PROJECT: undefined,
  246. GOOGLE_VERTEX_LOCATION: undefined,
  247. GOOGLE_CLOUD_LOCATION: undefined,
  248. VERTEX_LOCATION: undefined,
  249. },
  250. () =>
  251. Effect.gen(function* () {
  252. const catalog = yield* Catalog.Service
  253. yield* catalog.transform((catalog) =>
  254. catalog.provider.update(Provider.ID.make("google-vertex"), (provider) => {
  255. provider.package = Provider.aisdk("@ai-sdk/google-vertex")
  256. provider.settings = { ...provider.settings, project: "config-project" }
  257. }),
  258. )
  259. yield* addPlugin()
  260. const provider = required(yield* catalog.provider.get(Provider.ID.make("google-vertex")))
  261. expect(provider.settings?.project).toBe("config-project")
  262. expect(provider.settings?.location).toBe("us-central1")
  263. }),
  264. ),
  265. )
  266. it.effect("does not pass Google auth fetch to the native Vertex SDK", () =>
  267. withEnv(
  268. {
  269. GOOGLE_CLOUD_PROJECT: "env-project",
  270. GOOGLE_VERTEX_LOCATION: "env-location",
  271. },
  272. () =>
  273. Effect.gen(function* () {
  274. vertexOptions.length = 0
  275. const plugin = yield* Plugin.Service
  276. const aisdk = yield* AISDK.Service
  277. yield* addPlugin()
  278. yield* aisdk.runSDK({
  279. model: Model.Info.make({
  280. ...Model.Info.default(Provider.ID.make("google-vertex"), Model.ID.make("gemini")),
  281. modelID: Model.ID.make("gemini"),
  282. package: "aisdk:@ai-sdk/google-vertex",
  283. }),
  284. package: "@ai-sdk/google-vertex",
  285. options: { name: "google-vertex" },
  286. })
  287. expect(vertexOptions).toHaveLength(1)
  288. expect(vertexOptions[0].project).toBe("env-project")
  289. expect(vertexOptions[0].location).toBe("env-location")
  290. expect(vertexOptions[0].fetch).toBeUndefined()
  291. }),
  292. ),
  293. )
  294. it.effect("keeps Google auth fetch for OpenAI-compatible Vertex endpoints", () =>
  295. Effect.gen(function* () {
  296. googleAuthOptions.length = 0
  297. const fetchCalls: { input: Parameters<typeof fetch>[0]; init?: RequestInit }[] = []
  298. const plugin = yield* Plugin.Service
  299. const aisdk = yield* AISDK.Service
  300. yield* addPlugin()
  301. yield* aisdk.hook.sdk((evt) =>
  302. Effect.promise(async () => {
  303. if (evt.model.providerID !== "google-vertex") return
  304. if (evt.package !== "@ai-sdk/openai-compatible") return
  305. expect(typeof evt.options.fetch).toBe("function")
  306. await evt.options.fetch("https://vertex.example", {
  307. headers: { "x-test": "1" },
  308. })
  309. }),
  310. )
  311. const originalFetch = fetch
  312. ;(globalThis as typeof globalThis & { fetch: typeof fetch }).fetch = (async (
  313. input: Parameters<typeof fetch>[0],
  314. init?: RequestInit,
  315. ) => {
  316. fetchCalls.push({ input, init })
  317. return new Response("ok")
  318. }) as typeof fetch
  319. yield* Effect.acquireUseRelease(
  320. Effect.void,
  321. () =>
  322. aisdk.runSDK({
  323. model: Model.Info.make({
  324. ...Model.Info.default(Provider.ID.make("google-vertex"), Model.ID.make("gemini")),
  325. modelID: Model.ID.make("gemini"),
  326. package: "aisdk:@ai-sdk/openai-compatible",
  327. }),
  328. package: "@ai-sdk/openai-compatible",
  329. options: { name: "google-vertex" },
  330. }),
  331. () =>
  332. Effect.sync(() => {
  333. ;(globalThis as typeof globalThis & { fetch: typeof fetch }).fetch = originalFetch
  334. }),
  335. )
  336. const vertexCalls = fetchCalls.filter((call) => call.input === "https://vertex.example")
  337. expect(vertexCalls).toHaveLength(1)
  338. expect(googleAuthOptions).toEqual([{ scopes: ["https://www.googleapis.com/auth/cloud-platform"] }])
  339. expect(new Headers(vertexCalls[0].init?.headers).get("authorization")).toBe("Bearer vertex-token")
  340. expect(new Headers(vertexCalls[0].init?.headers).get("x-test")).toBe("1")
  341. }),
  342. )
  343. it.effect("trims model IDs before selecting language models", () =>
  344. Effect.gen(function* () {
  345. const plugin = yield* Plugin.Service
  346. const aisdk = yield* AISDK.Service
  347. const calls: string[] = []
  348. yield* addPlugin()
  349. yield* aisdk.runLanguage({
  350. model: Model.Info.make({
  351. ...Model.Info.default(Provider.ID.make("google-vertex"), Model.ID.make(" gemini-2.5-pro ")),
  352. modelID: Model.ID.make(" gemini-2.5-pro "),
  353. package: "aisdk:test-provider",
  354. }),
  355. sdk: { languageModel: fakeSelectorSdk(calls).languageModel },
  356. options: {},
  357. })
  358. expect(calls).toEqual(["languageModel:gemini-2.5-pro"])
  359. }),
  360. )
  361. })