diff --git a/src/plugin-sdk/provider-stream.test.ts b/src/plugin-sdk/provider-stream.test.ts index 05c69d01106..9bfceebd7f1 100644 --- a/src/plugin-sdk/provider-stream.test.ts +++ b/src/plugin-sdk/provider-stream.test.ts @@ -1,5 +1,6 @@ import type { StreamFn } from "@mariozechner/pi-agent-core"; import { describe, expect, it } from "vitest"; +import { composeProviderStreamWrappers as composeProviderStreamWrappersShared } from "./provider-stream-shared.js"; import { buildProviderStreamFamilyHooks, composeProviderStreamWrappers, @@ -24,6 +25,10 @@ function requireStreamFn(streamFn: StreamFn | null | undefined) { } describe("composeProviderStreamWrappers", () => { + it("re-exports the shared wrapper composer", () => { + expect(composeProviderStreamWrappers).toBe(composeProviderStreamWrappersShared); + }); + it("applies wrappers left to right", async () => { const order: string[] = []; const baseStreamFn: StreamFn = (_model, _context, _options) => { diff --git a/src/plugin-sdk/provider-stream.ts b/src/plugin-sdk/provider-stream.ts index d057843e138..82b07d9bbba 100644 --- a/src/plugin-sdk/provider-stream.ts +++ b/src/plugin-sdk/provider-stream.ts @@ -1,4 +1,3 @@ -import type { StreamFn } from "@mariozechner/pi-agent-core"; import { createGoogleThinkingPayloadWrapper, sanitizeGoogleThinkingPayload, @@ -32,22 +31,10 @@ import { } from "../agents/pi-embedded-runner/zai-stream-wrappers.js"; import type { ProviderPlugin } from "../plugins/types.js"; import type { ProviderWrapStreamFnContext } from "./plugin-entry.js"; - -export type ProviderStreamWrapperFactory = - | ((streamFn: StreamFn | undefined) => StreamFn | undefined) - | null - | undefined - | false; - -export function composeProviderStreamWrappers( - baseStreamFn: StreamFn | undefined, - ...wrappers: ProviderStreamWrapperFactory[] -): StreamFn | undefined { - return wrappers.reduce( - (streamFn, wrapper) => (wrapper ? wrapper(streamFn) : streamFn), - baseStreamFn, - ); -} +export { + composeProviderStreamWrappers, + type ProviderStreamWrapperFactory, +} from "./provider-stream-shared.js"; export type ProviderStreamFamily = | "google-thinking"