provider-nvidia.test.ts 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. import { describe, expect } from "bun:test"
  2. import { Effect } from "effect"
  3. import { Catalog } from "@opencode-ai/core/catalog"
  4. import { PluginV2 } from "@opencode-ai/core/plugin"
  5. import { ProviderPlugins } from "@opencode-ai/core/plugin/provider"
  6. import { NvidiaPlugin } from "@opencode-ai/core/plugin/provider/nvidia"
  7. import { ProviderV2 } from "@opencode-ai/core/provider"
  8. import { expectPluginRegistered, it, provider } from "./provider-helper"
  9. describe("NvidiaPlugin", () => {
  10. it.effect("is registered so legacy referer headers can be applied", () =>
  11. Effect.sync(() =>
  12. expectPluginRegistered(
  13. ProviderPlugins.map((item) => item.id),
  14. "nvidia",
  15. ),
  16. ),
  17. )
  18. it.effect("applies NVIDIA tracking headers only to nvidia", () =>
  19. Effect.gen(function* () {
  20. const plugin = yield* PluginV2.Service
  21. const catalog = yield* Catalog.Service
  22. yield* plugin.add(NvidiaPlugin)
  23. const load = yield* catalog.loader()
  24. yield* load((catalog) => {
  25. const nvidia = provider("nvidia", {
  26. endpoint: { type: "aisdk", package: "@ai-sdk/openai-compatible", url: "https://integrate.api.nvidia.com/v1" },
  27. options: { headers: { Existing: "value" }, body: {}, aisdk: { provider: {}, request: {} } },
  28. })
  29. catalog.provider.update(nvidia.id, (draft) => {
  30. draft.endpoint = nvidia.endpoint
  31. draft.options = nvidia.options
  32. })
  33. catalog.provider.update(provider("openrouter").id, () => {})
  34. })
  35. expect((yield* catalog.provider.get(ProviderV2.ID.make("nvidia"))).options.headers).toEqual({
  36. Existing: "value",
  37. "HTTP-Referer": "https://opencode.ai/",
  38. "X-Title": "opencode",
  39. "X-BILLING-INVOKE-ORIGIN": "OpenCode",
  40. })
  41. expect((yield* catalog.provider.get(ProviderV2.ID.openrouter)).options.headers).toEqual({})
  42. }),
  43. )
  44. it.effect("adds billing origin for custom NVIDIA endpoints", () =>
  45. Effect.gen(function* () {
  46. const plugin = yield* PluginV2.Service
  47. const catalog = yield* Catalog.Service
  48. yield* plugin.add(NvidiaPlugin)
  49. const load = yield* catalog.loader()
  50. yield* load((catalog) => {
  51. const item = provider("nvidia", {
  52. endpoint: { type: "aisdk", package: "@ai-sdk/openai-compatible", url: "https://integrate.api.nvidia.com/v1" },
  53. options: { headers: {}, body: {}, aisdk: { provider: {}, request: {} } },
  54. })
  55. catalog.provider.update(item.id, (draft) => {
  56. draft.endpoint = item.endpoint
  57. draft.options = item.options
  58. })
  59. })
  60. expect((yield* catalog.provider.get(ProviderV2.ID.make("nvidia"))).options.headers).toEqual({
  61. "HTTP-Referer": "https://opencode.ai/",
  62. "X-Title": "opencode",
  63. "X-BILLING-INVOKE-ORIGIN": "OpenCode",
  64. })
  65. }),
  66. )
  67. it.effect("preserves an explicit NVIDIA billing origin header", () =>
  68. Effect.gen(function* () {
  69. const plugin = yield* PluginV2.Service
  70. const catalog = yield* Catalog.Service
  71. yield* plugin.add(NvidiaPlugin)
  72. const load = yield* catalog.loader()
  73. yield* load((catalog) => {
  74. const item = provider("nvidia", {
  75. endpoint: { type: "aisdk", package: "@ai-sdk/openai-compatible", url: "https://integrate.api.nvidia.com/v1" },
  76. options: {
  77. headers: { "X-BILLING-INVOKE-ORIGIN": "CustomOrigin" },
  78. body: {},
  79. aisdk: { provider: { baseURL: "https://integrate.api.nvidia.com/v1" }, request: {} },
  80. },
  81. })
  82. catalog.provider.update(item.id, (draft) => {
  83. draft.endpoint = item.endpoint
  84. draft.options = item.options
  85. })
  86. })
  87. expect((yield* catalog.provider.get(ProviderV2.ID.make("nvidia"))).options.headers).toEqual({
  88. "HTTP-Referer": "https://opencode.ai/",
  89. "X-Title": "opencode",
  90. "X-BILLING-INVOKE-ORIGIN": "CustomOrigin",
  91. })
  92. }),
  93. )
  94. })