provider-selector-model-mismatch-binding.test.ts 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192
  1. import { beforeEach, describe, expect, test, vi } from "vitest";
  2. import type { Provider } from "@/types/provider";
  3. const circuitBreakerMocks = vi.hoisted(() => ({
  4. isCircuitOpen: vi.fn(async () => false),
  5. getCircuitState: vi.fn(() => "closed"),
  6. }));
  7. vi.mock("@/lib/circuit-breaker", () => circuitBreakerMocks);
  8. const vendorTypeCircuitMocks = vi.hoisted(() => ({
  9. isVendorTypeCircuitOpen: vi.fn(async () => false),
  10. }));
  11. vi.mock("@/lib/vendor-type-circuit-breaker", () => vendorTypeCircuitMocks);
  12. const sessionManagerMocks = vi.hoisted(() => ({
  13. SessionManager: {
  14. getSessionProvider: vi.fn(async () => null as number | null),
  15. clearSessionProvider: vi.fn(async () => undefined),
  16. },
  17. }));
  18. vi.mock("@/lib/session-manager", () => sessionManagerMocks);
  19. const providerRepositoryMocks = vi.hoisted(() => ({
  20. findProviderById: vi.fn(async () => null as Provider | null),
  21. findAllProviders: vi.fn(async () => [] as Provider[]),
  22. }));
  23. vi.mock("@/repository/provider", () => providerRepositoryMocks);
  24. const rateLimitMocks = vi.hoisted(() => ({
  25. RateLimitService: {
  26. checkCostLimitsWithLease: vi.fn(async () => ({ allowed: true })),
  27. checkTotalCostLimit: vi.fn(async () => ({ allowed: true, current: 0 })),
  28. },
  29. }));
  30. vi.mock("@/lib/rate-limit", () => rateLimitMocks);
  31. beforeEach(() => {
  32. vi.resetAllMocks();
  33. });
  34. function createHaikuOnlyProvider(): Provider {
  35. return {
  36. id: 78,
  37. name: "zhipu_Haiku",
  38. isEnabled: true,
  39. providerType: "claude",
  40. groupTag: null,
  41. weight: 1,
  42. priority: 1,
  43. costMultiplier: 1,
  44. allowedModels: ["claude-haiku-4-5-20251001", "claude-haiku-4-5"],
  45. providerVendorId: null,
  46. limit5hUsd: null,
  47. limitDailyUsd: null,
  48. dailyResetMode: "fixed",
  49. dailyResetTime: "00:00",
  50. limitWeeklyUsd: null,
  51. limitMonthlyUsd: null,
  52. limitTotalUsd: null,
  53. totalCostResetAt: null,
  54. limitConcurrentSessions: 0,
  55. } as unknown as Provider;
  56. }
  57. function createOpusProvider(): Provider {
  58. return {
  59. id: 94,
  60. name: "yescode_team",
  61. isEnabled: true,
  62. providerType: "claude",
  63. groupTag: null,
  64. weight: 1,
  65. priority: 0,
  66. costMultiplier: 1,
  67. allowedModels: null, // supports all claude models
  68. providerVendorId: null,
  69. limit5hUsd: null,
  70. limitDailyUsd: null,
  71. dailyResetMode: "fixed",
  72. dailyResetTime: "00:00",
  73. limitWeeklyUsd: null,
  74. limitMonthlyUsd: null,
  75. limitTotalUsd: null,
  76. totalCostResetAt: null,
  77. limitConcurrentSessions: 0,
  78. } as unknown as Provider;
  79. }
  80. describe("findReusable - model mismatch clears stale binding", () => {
  81. test("should clear stale binding when bound provider does not support requested model", async () => {
  82. const { ProxyProviderResolver } = await import("@/app/v1/_lib/proxy/provider-selector");
  83. // Session bound to haiku-only provider
  84. sessionManagerMocks.SessionManager.getSessionProvider.mockResolvedValueOnce(78);
  85. providerRepositoryMocks.findProviderById.mockResolvedValueOnce(createHaikuOnlyProvider());
  86. const session = {
  87. sessionId: "4c25cf92",
  88. shouldReuseProvider: () => true,
  89. getOriginalModel: () => "claude-opus-4-6",
  90. authState: null,
  91. getCurrentModel: () => null,
  92. } as any;
  93. const result = await (ProxyProviderResolver as any).findReusable(session);
  94. expect(result).toBeNull();
  95. // Key assertion: clearSessionProvider should have been called
  96. expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith(
  97. "4c25cf92"
  98. );
  99. });
  100. test("should NOT clear binding when bound provider supports requested model", async () => {
  101. const { ProxyProviderResolver } = await import("@/app/v1/_lib/proxy/provider-selector");
  102. // Session bound to provider that supports all claude models
  103. sessionManagerMocks.SessionManager.getSessionProvider.mockResolvedValueOnce(94);
  104. providerRepositoryMocks.findProviderById.mockResolvedValueOnce(createOpusProvider());
  105. rateLimitMocks.RateLimitService.checkCostLimitsWithLease.mockResolvedValueOnce({
  106. allowed: true,
  107. });
  108. rateLimitMocks.RateLimitService.checkTotalCostLimit.mockResolvedValueOnce({
  109. allowed: true,
  110. current: 0,
  111. });
  112. const session = {
  113. sessionId: "sess_ok",
  114. shouldReuseProvider: () => true,
  115. getOriginalModel: () => "claude-opus-4-6",
  116. authState: null,
  117. getCurrentModel: () => null,
  118. } as any;
  119. const result = await (ProxyProviderResolver as any).findReusable(session);
  120. // Should return the provider (model matches)
  121. expect(result).not.toBeNull();
  122. expect(result?.id).toBe(94);
  123. // clearSessionProvider should NOT have been called
  124. expect(sessionManagerMocks.SessionManager.clearSessionProvider).not.toHaveBeenCalled();
  125. });
  126. test("should NOT clear binding when shouldReuseProvider returns false", async () => {
  127. const { ProxyProviderResolver } = await import("@/app/v1/_lib/proxy/provider-selector");
  128. const session = {
  129. sessionId: "sess_short",
  130. shouldReuseProvider: () => false,
  131. getOriginalModel: () => "claude-opus-4-6",
  132. authState: null,
  133. } as any;
  134. const result = await (ProxyProviderResolver as any).findReusable(session);
  135. expect(result).toBeNull();
  136. // Should not even reach the model check, so no clear
  137. expect(sessionManagerMocks.SessionManager.clearSessionProvider).not.toHaveBeenCalled();
  138. expect(sessionManagerMocks.SessionManager.getSessionProvider).not.toHaveBeenCalled();
  139. });
  140. test("should clear binding for haiku-only provider when requesting haiku-4-5 variant not in allowlist", async () => {
  141. const { ProxyProviderResolver } = await import("@/app/v1/_lib/proxy/provider-selector");
  142. sessionManagerMocks.SessionManager.getSessionProvider.mockResolvedValueOnce(78);
  143. const provider = createHaikuOnlyProvider();
  144. // Restrictive allowlist - only allows specific variant
  145. provider.allowedModels = ["claude-haiku-4-5-20251001"];
  146. providerRepositoryMocks.findProviderById.mockResolvedValueOnce(provider);
  147. const session = {
  148. sessionId: "sess_variant",
  149. shouldReuseProvider: () => true,
  150. getOriginalModel: () => "claude-sonnet-4-5-20250929",
  151. authState: null,
  152. getCurrentModel: () => null,
  153. } as any;
  154. const result = await (ProxyProviderResolver as any).findReusable(session);
  155. expect(result).toBeNull();
  156. expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith(
  157. "sess_variant"
  158. );
  159. });
  160. });