provider-google-vertex.test.ts 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300
  1. import { describe, expect, mock } from "bun:test"
  2. import { Effect } from "effect"
  3. import { PluginV2 } from "@opencode-ai/core/plugin"
  4. import { GoogleVertexPlugin } from "@opencode-ai/core/plugin/provider/google-vertex"
  5. import { fakeSelectorSdk, it, model, provider, withEnv } from "./provider-helper"
  6. const vertexOptions: Record<string, any>[] = []
  7. void mock.module("@ai-sdk/google-vertex", () => ({
  8. createVertex: (options: Record<string, any>) => {
  9. vertexOptions.push(options)
  10. return {
  11. languageModel: (modelID: string) => ({ modelID, provider: "google-vertex", specificationVersion: "v3" }),
  12. }
  13. },
  14. }))
  15. void mock.module("google-auth-library", () => ({
  16. GoogleAuth: class {
  17. async getApplicationDefault() {
  18. return {
  19. credential: {
  20. async getAccessToken() {
  21. return { token: "vertex-token" }
  22. },
  23. },
  24. }
  25. }
  26. },
  27. }))
  28. describe("GoogleVertexPlugin", () => {
  29. it.effect("resolves project and location from env using legacy precedence", () =>
  30. withEnv(
  31. {
  32. GOOGLE_CLOUD_PROJECT: "google-cloud-project",
  33. GCP_PROJECT: "gcp-project",
  34. GCLOUD_PROJECT: "gcloud-project",
  35. GOOGLE_VERTEX_LOCATION: "google-vertex-location",
  36. GOOGLE_CLOUD_LOCATION: "google-cloud-location",
  37. VERTEX_LOCATION: "vertex-location",
  38. },
  39. () =>
  40. Effect.gen(function* () {
  41. const plugin = yield* PluginV2.Service
  42. yield* plugin.add(GoogleVertexPlugin)
  43. const result = yield* plugin.trigger(
  44. "provider.update",
  45. {},
  46. {
  47. provider: provider("google-vertex", {
  48. endpoint: {
  49. type: "aisdk",
  50. package: "@ai-sdk/openai-compatible",
  51. url: "https://${GOOGLE_VERTEX_ENDPOINT}/v1/projects/${GOOGLE_VERTEX_PROJECT}/locations/${GOOGLE_VERTEX_LOCATION}",
  52. },
  53. }),
  54. cancel: false,
  55. },
  56. )
  57. expect(result.provider.options.aisdk.provider.project).toBe("google-cloud-project")
  58. expect(result.provider.options.aisdk.provider.location).toBe("google-vertex-location")
  59. expect(result.provider.endpoint).toEqual({
  60. type: "aisdk",
  61. package: "@ai-sdk/openai-compatible",
  62. url: "https://google-vertex-location-aiplatform.googleapis.com/v1/projects/google-cloud-project/locations/google-vertex-location",
  63. })
  64. }),
  65. ),
  66. )
  67. it.effect("resolves the advertised GOOGLE_VERTEX_PROJECT env for provider updates and SDKs", () =>
  68. withEnv(
  69. {
  70. GOOGLE_VERTEX_PROJECT: "vertex-project",
  71. GOOGLE_CLOUD_PROJECT: undefined,
  72. GCP_PROJECT: undefined,
  73. GCLOUD_PROJECT: undefined,
  74. GOOGLE_VERTEX_LOCATION: "europe-west4",
  75. GOOGLE_CLOUD_LOCATION: undefined,
  76. VERTEX_LOCATION: undefined,
  77. },
  78. () =>
  79. Effect.gen(function* () {
  80. vertexOptions.length = 0
  81. const plugin = yield* PluginV2.Service
  82. yield* plugin.add(GoogleVertexPlugin)
  83. const updated = yield* plugin.trigger(
  84. "provider.update",
  85. {},
  86. {
  87. provider: provider("google-vertex", {
  88. endpoint: {
  89. type: "aisdk",
  90. package: "@ai-sdk/openai-compatible",
  91. url: "https://${GOOGLE_VERTEX_ENDPOINT}/v1/projects/${GOOGLE_VERTEX_PROJECT}/locations/${GOOGLE_VERTEX_LOCATION}",
  92. },
  93. }),
  94. cancel: false,
  95. },
  96. )
  97. yield* plugin.trigger(
  98. "aisdk.sdk",
  99. {
  100. model: model("google-vertex", "gemini", {
  101. endpoint: { type: "aisdk", package: "@ai-sdk/google-vertex" },
  102. }),
  103. package: "@ai-sdk/google-vertex",
  104. options: { name: "google-vertex" },
  105. },
  106. {},
  107. )
  108. expect(updated.provider.options.aisdk.provider.project).toBe("vertex-project")
  109. expect(updated.provider.endpoint).toEqual({
  110. type: "aisdk",
  111. package: "@ai-sdk/openai-compatible",
  112. url: "https://europe-west4-aiplatform.googleapis.com/v1/projects/vertex-project/locations/europe-west4",
  113. })
  114. expect(vertexOptions[0].project).toBe("vertex-project")
  115. expect(vertexOptions[0].location).toBe("europe-west4")
  116. }),
  117. ),
  118. )
  119. it.effect("keeps configured project and location over env and uses global endpoint", () =>
  120. withEnv(
  121. {
  122. GOOGLE_CLOUD_PROJECT: "env-project",
  123. GCP_PROJECT: "env-gcp-project",
  124. GCLOUD_PROJECT: "env-gcloud-project",
  125. GOOGLE_VERTEX_LOCATION: "env-location",
  126. GOOGLE_CLOUD_LOCATION: "env-google-cloud-location",
  127. VERTEX_LOCATION: "env-vertex-location",
  128. },
  129. () =>
  130. Effect.gen(function* () {
  131. const plugin = yield* PluginV2.Service
  132. yield* plugin.add(GoogleVertexPlugin)
  133. const result = yield* plugin.trigger(
  134. "provider.update",
  135. {},
  136. {
  137. provider: provider("google-vertex", {
  138. endpoint: {
  139. type: "aisdk",
  140. package: "@ai-sdk/openai-compatible",
  141. url: "https://${GOOGLE_VERTEX_ENDPOINT}/v1/projects/${GOOGLE_VERTEX_PROJECT}/locations/${GOOGLE_VERTEX_LOCATION}",
  142. },
  143. options: {
  144. headers: {},
  145. body: {},
  146. aisdk: { provider: { project: "config-project", location: "global" }, request: {} },
  147. },
  148. }),
  149. cancel: false,
  150. },
  151. )
  152. expect(result.provider.options.aisdk.provider.project).toBe("config-project")
  153. expect(result.provider.options.aisdk.provider.location).toBe("global")
  154. expect(result.provider.endpoint).toEqual({
  155. type: "aisdk",
  156. package: "@ai-sdk/openai-compatible",
  157. url: "https://aiplatform.googleapis.com/v1/projects/config-project/locations/global",
  158. })
  159. }),
  160. ),
  161. )
  162. it.effect("defaults location to us-central1 when only project is configured", () =>
  163. withEnv(
  164. {
  165. GOOGLE_CLOUD_PROJECT: undefined,
  166. GCP_PROJECT: undefined,
  167. GCLOUD_PROJECT: undefined,
  168. GOOGLE_VERTEX_LOCATION: undefined,
  169. GOOGLE_CLOUD_LOCATION: undefined,
  170. VERTEX_LOCATION: undefined,
  171. },
  172. () =>
  173. Effect.gen(function* () {
  174. const plugin = yield* PluginV2.Service
  175. yield* plugin.add(GoogleVertexPlugin)
  176. const result = yield* plugin.trigger(
  177. "provider.update",
  178. {},
  179. {
  180. provider: provider("google-vertex", {
  181. options: { headers: {}, body: {}, aisdk: { provider: { project: "config-project" }, request: {} } },
  182. }),
  183. cancel: false,
  184. },
  185. )
  186. expect(result.provider.options.aisdk.provider.project).toBe("config-project")
  187. expect(result.provider.options.aisdk.provider.location).toBe("us-central1")
  188. }),
  189. ),
  190. )
  191. it.effect("does not pass Google auth fetch to the native Vertex SDK", () =>
  192. withEnv(
  193. {
  194. GOOGLE_CLOUD_PROJECT: "env-project",
  195. GOOGLE_VERTEX_LOCATION: "env-location",
  196. },
  197. () =>
  198. Effect.gen(function* () {
  199. vertexOptions.length = 0
  200. const plugin = yield* PluginV2.Service
  201. yield* plugin.add(GoogleVertexPlugin)
  202. yield* plugin.trigger(
  203. "aisdk.sdk",
  204. {
  205. model: model("google-vertex", "gemini", {
  206. endpoint: { type: "aisdk", package: "@ai-sdk/google-vertex" },
  207. }),
  208. package: "@ai-sdk/google-vertex",
  209. options: { name: "google-vertex" },
  210. },
  211. {},
  212. )
  213. expect(vertexOptions).toHaveLength(1)
  214. expect(vertexOptions[0].project).toBe("env-project")
  215. expect(vertexOptions[0].location).toBe("env-location")
  216. expect(vertexOptions[0].fetch).toBeUndefined()
  217. }),
  218. ),
  219. )
  220. it.effect("keeps Google auth fetch for OpenAI-compatible Vertex endpoints", () =>
  221. Effect.gen(function* () {
  222. const fetchCalls: { input: Parameters<typeof fetch>[0]; init?: RequestInit }[] = []
  223. const plugin = yield* PluginV2.Service
  224. yield* plugin.add(GoogleVertexPlugin)
  225. yield* plugin.add({
  226. id: PluginV2.ID.make("capture-openai-compatible"),
  227. effect: Effect.succeed({
  228. "aisdk.sdk": (evt) =>
  229. Effect.promise(async () => {
  230. if (evt.model.providerID !== "google-vertex") return
  231. if (evt.package !== "@ai-sdk/openai-compatible") return
  232. expect(typeof evt.options.fetch).toBe("function")
  233. await evt.options.fetch("https://vertex.example", {
  234. headers: { "x-test": "1" },
  235. })
  236. }),
  237. }),
  238. })
  239. const originalFetch = fetch
  240. ;(globalThis as typeof globalThis & { fetch: typeof fetch }).fetch = (async (
  241. input: Parameters<typeof fetch>[0],
  242. init?: RequestInit,
  243. ) => {
  244. fetchCalls.push({ input, init })
  245. return new Response("ok")
  246. }) as typeof fetch
  247. yield* Effect.acquireUseRelease(
  248. Effect.void,
  249. () =>
  250. plugin.trigger(
  251. "aisdk.sdk",
  252. {
  253. model: model("google-vertex", "gemini", {
  254. endpoint: { type: "aisdk", package: "@ai-sdk/openai-compatible" },
  255. }),
  256. package: "@ai-sdk/openai-compatible",
  257. options: { name: "google-vertex" },
  258. },
  259. {},
  260. ),
  261. () =>
  262. Effect.sync(() => {
  263. ;(globalThis as typeof globalThis & { fetch: typeof fetch }).fetch = originalFetch
  264. }),
  265. )
  266. expect(fetchCalls).toHaveLength(1)
  267. expect(fetchCalls[0].input).toBe("https://vertex.example")
  268. expect(new Headers(fetchCalls[0].init?.headers).get("authorization")).toBe("Bearer vertex-token")
  269. expect(new Headers(fetchCalls[0].init?.headers).get("x-test")).toBe("1")
  270. }),
  271. )
  272. it.effect("trims model IDs before selecting language models", () =>
  273. Effect.gen(function* () {
  274. const plugin = yield* PluginV2.Service
  275. const calls: string[] = []
  276. yield* plugin.add(GoogleVertexPlugin)
  277. yield* plugin.trigger(
  278. "aisdk.language",
  279. {
  280. model: model("google-vertex", " gemini-2.5-pro "),
  281. sdk: { languageModel: fakeSelectorSdk(calls).languageModel },
  282. options: {},
  283. },
  284. {},
  285. )
  286. expect(calls).toEqual(["languageModel:gemini-2.5-pro"])
  287. }),
  288. )
  289. })