diff --git a/extensions/anthropic/stream-wrappers.test.ts b/extensions/anthropic/stream-wrappers.test.ts index b82c74d0128..04481f2964e 100644 --- a/extensions/anthropic/stream-wrappers.test.ts +++ b/extensions/anthropic/stream-wrappers.test.ts @@ -26,6 +26,59 @@ function runWrapper(apiKey: string | undefined): Record | undefi return captured.headers; } +function createPayloadCapturingBaseStream(captured: { + headers?: Record; + payload?: Record; +}): StreamFn { + return (model, _context, options) => { + captured.headers = options?.headers; + const payload = {} as Record; + options?.onPayload?.(payload as never, model as never); + captured.payload = payload; + return {} as never; + }; +} + +function runComposedAnthropicProviderStream(apiKey: string) { + const captured: { headers?: Record; payload?: Record } = {}; + const wrapped = wrapAnthropicProviderStream({ + streamFn: createPayloadCapturingBaseStream(captured), + modelId: "claude-sonnet-4-6", + extraParams: { context1m: true, serviceTier: "auto" }, + } as never); + + void wrapped?.( + { provider: "anthropic", api: "anthropic-messages", id: "claude-sonnet-4-6" } as never, + {} as never, + { apiKey } as never, + ); + return captured; +} + +function runPayloadWrapper( + params: { + apiKey?: string; + provider?: string; + api?: string; + baseUrl?: string; + }, + createWrapper: (base: StreamFn) => StreamFn, +): Record | undefined { + const captured: { payload?: Record } = {}; + const wrapper = createWrapper(createPayloadCapturingBaseStream(captured)); + void wrapper( + { + provider: params.provider ?? "anthropic", + api: params.api ?? "anthropic-messages", + baseUrl: params.baseUrl, + id: "claude-sonnet-4-6", + } as never, + {} as never, + { apiKey: params.apiKey } as never, + ); + return captured.payload; +} + describe("anthropic stream wrappers", () => { afterEach(() => { vi.restoreAllMocks(); @@ -49,54 +102,14 @@ describe("anthropic stream wrappers", () => { }); it("skips service_tier for OAuth token in composed stream chain", () => { - const captured: { headers?: Record; payload?: Record } = {}; - const base: StreamFn = (model, _context, options) => { - captured.headers = options?.headers; - const payload = {} as Record; - options?.onPayload?.(payload as never, model as never); - captured.payload = payload; - return {} as never; - }; - - const wrapped = wrapAnthropicProviderStream({ - streamFn: base, - modelId: "claude-sonnet-4-6", - extraParams: { context1m: true, serviceTier: "auto" }, - } as never); - - void wrapped?.( - { provider: "anthropic", api: "anthropic-messages", id: "claude-sonnet-4-6" } as never, - {} as never, - { apiKey: "sk-ant-oat01-oauth-token" } as never, - ); - + const captured = runComposedAnthropicProviderStream("sk-ant-oat01-oauth-token"); expect(captured.headers?.["anthropic-beta"]).toContain(OAUTH_BETA); expect(captured.headers?.["anthropic-beta"]).not.toContain(CONTEXT_1M_BETA); expect(captured.payload?.service_tier).toBeUndefined(); }); it("composes the anthropic provider stream chain from extra params", () => { - const captured: { headers?: Record; payload?: Record } = {}; - const base: StreamFn = (model, _context, options) => { - captured.headers = options?.headers; - const payload = {} as Record; - options?.onPayload?.(payload as never, model as never); - captured.payload = payload; - return {} as never; - }; - - const wrapped = wrapAnthropicProviderStream({ - streamFn: base, - modelId: "claude-sonnet-4-6", - extraParams: { context1m: true, serviceTier: "auto" }, - } as never); - - void wrapped?.( - { provider: "anthropic", api: "anthropic-messages", id: "claude-sonnet-4-6" } as never, - {} as never, - { apiKey: "sk-ant-api-123" } as never, - ); - + const captured = runComposedAnthropicProviderStream("sk-ant-api-123"); expect(captured.headers?.["anthropic-beta"]).toContain(CONTEXT_1M_BETA); expect(captured.payload).toMatchObject({ service_tier: "auto" }); }); @@ -110,28 +123,9 @@ describe("createAnthropicFastModeWrapper", () => { baseUrl?: string; enabled?: boolean; }): Record | undefined { - const captured: { payload?: Record } = {}; - const base: StreamFn = (_model, _context, options) => { - if (options?.onPayload) { - const payload: Record = {}; - options.onPayload(payload, _model); - captured.payload = payload; - } - return {} as never; - }; - - const wrapper = createAnthropicFastModeWrapper(base, params.enabled ?? true); - void wrapper( - { - provider: params.provider ?? "anthropic", - api: params.api ?? "anthropic-messages", - baseUrl: params.baseUrl, - id: "claude-sonnet-4-6", - } as never, - {} as never, - { apiKey: params.apiKey } as never, + return runPayloadWrapper(params, (base) => + createAnthropicFastModeWrapper(base, params.enabled ?? true), ); - return captured.payload; } it("does not inject service_tier for OAuth token", () => { @@ -166,27 +160,9 @@ describe("createAnthropicServiceTierWrapper", () => { api?: string; serviceTier?: "auto" | "standard_only"; }): Record | undefined { - const captured: { payload?: Record } = {}; - const base: StreamFn = (_model, _context, options) => { - if (options?.onPayload) { - const payload: Record = {}; - options.onPayload(payload, _model); - captured.payload = payload; - } - return {} as never; - }; - - const wrapper = createAnthropicServiceTierWrapper(base, params.serviceTier ?? "auto"); - void wrapper( - { - provider: params.provider ?? "anthropic", - api: params.api ?? "anthropic-messages", - id: "claude-sonnet-4-6", - } as never, - {} as never, - { apiKey: params.apiKey } as never, + return runPayloadWrapper(params, (base) => + createAnthropicServiceTierWrapper(base, params.serviceTier ?? "auto"), ); - return captured.payload; } it("does not inject service_tier for OAuth token", () => {