vultr/openclaw
1
0
Fork 0
mirror of https://github.com/openclaw/openclaw.git synced 2026-04-23 07:01:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(providers): add stream family hooks

Browse Source
This commit is contained in:
Vincent Koc
2026-04-04 20:19:43 +09:00
parent 1037af01ad
commit bc648ac8e6
7 changed files with 162 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
extensions/google/gemini-cli-provider.ts
View File

@@ -2,11 +2,10 @@ import type {
OpenClawPluginApi,
ProviderAuthContext,
ProviderFetchUsageSnapshotContext,
ProviderWrapStreamFnContext,
} from "openclaw/plugin-sdk/plugin-entry";
import { buildOauthProviderAuthResult } from "openclaw/plugin-sdk/provider-auth-result";
import { buildProviderReplayFamilyHooks } from "openclaw/plugin-sdk/provider-model-shared";
import { createGoogleThinkingPayloadWrapper } from "openclaw/plugin-sdk/provider-stream";
import { buildProviderStreamFamilyHooks } from "openclaw/plugin-sdk/provider-stream";
import { buildProviderToolCompatFamilyHooks } from "openclaw/plugin-sdk/provider-tools";
import { fetchGeminiUsage } from "openclaw/plugin-sdk/provider-usage";
import { formatGoogleOauthApiKey, parseGoogleUsageToken } from "./oauth-token-shared.js";
@@ -24,8 +23,7 @@ const ENV_VARS = [
const GOOGLE_GEMINI_CLI_PROVIDER_HOOKS = {
...buildProviderReplayFamilyHooks({ family: "google-gemini" }),
wrapStreamFn: (ctx: ProviderWrapStreamFnContext) =>
createGoogleThinkingPayloadWrapper(ctx.streamFn, ctx.thinkingLevel),
...buildProviderStreamFamilyHooks("google-thinking"),
...buildProviderToolCompatFamilyHooks("gemini"),
};

6
extensions/google/index.ts
View File

@@ -5,12 +5,11 @@ import {
type OpenClawPluginApi,
type ProviderAuthContext,
type ProviderFetchUsageSnapshotContext,
type ProviderWrapStreamFnContext,
} from "openclaw/plugin-sdk/plugin-entry";
import { createProviderApiKeyAuthMethod } from "openclaw/plugin-sdk/provider-auth-api-key";
import type { ProviderPlugin } from "openclaw/plugin-sdk/provider-model-shared";
import { buildProviderReplayFamilyHooks } from "openclaw/plugin-sdk/provider-model-shared";
import { createGoogleThinkingPayloadWrapper } from "openclaw/plugin-sdk/provider-stream";
import { buildProviderStreamFamilyHooks } from "openclaw/plugin-sdk/provider-stream";
import { buildProviderToolCompatFamilyHooks } from "openclaw/plugin-sdk/provider-tools";
import {
GOOGLE_GEMINI_DEFAULT_MODEL,
@@ -50,8 +49,7 @@ const GOOGLE_GEMINI_REPLAY_HOOKS = buildProviderReplayFamilyHooks({
});
const GOOGLE_GEMINI_PROVIDER_HOOKS = {
...GOOGLE_GEMINI_REPLAY_HOOKS,
wrapStreamFn: (ctx: ProviderWrapStreamFnContext) =>
createGoogleThinkingPayloadWrapper(ctx.streamFn, ctx.thinkingLevel),
...buildProviderStreamFamilyHooks("google-thinking"),
};
const GOOGLE_GEMINI_PROVIDER_HOOKS_WITH_TOOL_COMPAT = {
...GOOGLE_GEMINI_PROVIDER_HOOKS,

9
extensions/minimax/index.ts
View File

@@ -12,7 +12,7 @@ import {
import { buildOauthProviderAuthResult } from "openclaw/plugin-sdk/provider-auth";
import { createProviderApiKeyAuthMethod } from "openclaw/plugin-sdk/provider-auth-api-key";
import { buildProviderReplayFamilyHooks } from "openclaw/plugin-sdk/provider-model-shared";
import { createMinimaxFastModeWrapper } from "openclaw/plugin-sdk/provider-stream";
import { buildProviderStreamFamilyHooks } from "openclaw/plugin-sdk/provider-stream";
import { fetchMinimaxUsage } from "openclaw/plugin-sdk/provider-usage";
import { isMiniMaxModernModelId, MINIMAX_DEFAULT_MODEL_ID } from "./api.js";
import {
@@ -39,6 +39,7 @@ const HYBRID_ANTHROPIC_OPENAI_REPLAY_HOOKS = buildProviderReplayFamilyHooks({
family: "hybrid-anthropic-openai",
anthropicModelDropThinkingBlocks: true,
});
const MINIMAX_FAST_MODE_STREAM_HOOKS = buildProviderStreamFamilyHooks("minimax-fast-mode");
function resolveMinimaxReasoningOutputMode(): "native" {
// Keep MiniMax on native reasoning mode. Tagged enforcement previously
@@ -247,8 +248,7 @@ export default definePluginEntry({
return apiKey ? { token: apiKey } : null;
},
...HYBRID_ANTHROPIC_OPENAI_REPLAY_HOOKS,
wrapStreamFn: (ctx) =>
createMinimaxFastModeWrapper(ctx.streamFn, ctx.extraParams?.fastMode === true),
...MINIMAX_FAST_MODE_STREAM_HOOKS,
resolveReasoningOutputMode: () => resolveMinimaxReasoningOutputMode(),
isModernModelRef: ({ modelId }) => isMiniMaxModernModelId(modelId),
fetchUsageSnapshot: async (ctx) =>
@@ -300,8 +300,7 @@ export default definePluginEntry({
},
],
...HYBRID_ANTHROPIC_OPENAI_REPLAY_HOOKS,
wrapStreamFn: (ctx) =>
createMinimaxFastModeWrapper(ctx.streamFn, ctx.extraParams?.fastMode === true),
...MINIMAX_FAST_MODE_STREAM_HOOKS,
resolveReasoningOutputMode: () => resolveMinimaxReasoningOutputMode(),
isModernModelRef: ({ modelId }) => isMiniMaxModernModelId(modelId),
});

14
extensions/moonshot/index.ts
View File

@@ -1,9 +1,6 @@
import { defineSingleProviderPluginEntry } from "openclaw/plugin-sdk/provider-entry";
import { buildProviderReplayFamilyHooks } from "openclaw/plugin-sdk/provider-model-shared";
import {
createMoonshotThinkingWrapper,
resolveMoonshotThinkingType,
} from "openclaw/plugin-sdk/provider-moonshot";
import { buildProviderStreamFamilyHooks } from "openclaw/plugin-sdk/provider-stream";
import { applyMoonshotNativeStreamingUsageCompat } from "./api.js";
import { moonshotMediaUnderstandingProvider } from "./media-understanding-provider.js";
import {
@@ -18,6 +15,7 @@ const PROVIDER_ID = "moonshot";
const OPENAI_COMPATIBLE_REPLAY_HOOKS = buildProviderReplayFamilyHooks({
family: "openai-compatible",
});
const MOONSHOT_THINKING_STREAM_HOOKS = buildProviderStreamFamilyHooks("moonshot-thinking");
export default defineSingleProviderPluginEntry({
id: PROVIDER_ID,
@@ -63,13 +61,7 @@ export default defineSingleProviderPluginEntry({
applyNativeStreamingUsageCompat: ({ providerConfig }) =>
applyMoonshotNativeStreamingUsageCompat(providerConfig),
...OPENAI_COMPATIBLE_REPLAY_HOOKS,
wrapStreamFn: (ctx) => {
const thinkingType = resolveMoonshotThinkingType({
configuredThinking: ctx.extraParams?.thinking,
thinkingLevel: ctx.thinkingLevel,
});
return createMoonshotThinkingWrapper(ctx.streamFn, thinkingType);
},
...MOONSHOT_THINKING_STREAM_HOOKS,
},
register(api) {
api.registerMediaUnderstandingProvider(moonshotMediaUnderstandingProvider);

6
extensions/zai/index.ts
View File

@@ -20,7 +20,7 @@ import {
buildProviderReplayFamilyHooks,
normalizeModelCompat,
} from "openclaw/plugin-sdk/provider-model-shared";
import { createZaiToolStreamWrapper } from "openclaw/plugin-sdk/provider-stream";
import { buildProviderStreamFamilyHooks } from "openclaw/plugin-sdk/provider-stream";
import { fetchZaiUsage, resolveLegacyPiAgentAccessToken } from "openclaw/plugin-sdk/provider-usage";
import { detectZaiEndpoint, type ZaiEndpointId } from "./detect.js";
import { zaiMediaUnderstandingProvider } from "./media-understanding-provider.js";
@@ -33,6 +33,7 @@ const PROFILE_ID = "zai:default";
const OPENAI_COMPATIBLE_REPLAY_HOOKS = buildProviderReplayFamilyHooks({
family: "openai-compatible",
});
const ZAI_TOOL_STREAM_HOOKS = buildProviderStreamFamilyHooks("tool-stream-default-on");
function resolveGlm5ForwardCompatModel(
ctx: ProviderResolveDynamicModelContext,
@@ -288,8 +289,7 @@ export default definePluginEntry({
tool_stream: true,
};
},
wrapStreamFn: (ctx) =>
createZaiToolStreamWrapper(ctx.streamFn, ctx.extraParams?.tool_stream !== false),
...ZAI_TOOL_STREAM_HOOKS,
isBinaryThinking: () => true,
isModernModelRef: ({ modelId }) => {
const lower = modelId.trim().toLowerCase();

88
src/plugin-sdk/provider-stream.test.ts
View File

@@ -1,6 +1,9 @@
import type { StreamFn } from "@mariozechner/pi-agent-core";
import { describe, expect, it } from "vitest";
import { composeProviderStreamWrappers } from "./provider-stream.js";
import {
buildProviderStreamFamilyHooks,
composeProviderStreamWrappers,
} from "./provider-stream.js";
describe("composeProviderStreamWrappers", () => {
it("applies wrappers left to right", async () => {
@@ -33,3 +36,86 @@ describe("composeProviderStreamWrappers", () => {
expect(composeProviderStreamWrappers(baseStreamFn)).toBe(baseStreamFn);
});
});
describe("buildProviderStreamFamilyHooks", () => {
it("covers the stream family matrix", () => {
let capturedPayload: Record<string, unknown> | undefined;
let capturedModelId: string | undefined;
const baseStreamFn: StreamFn = (model, _context, options) => {
capturedModelId = String(model.id);
const payload = { config: { thinkingConfig: { thinkingBudget: -1 } } } as Record<
string,
unknown
>;
options?.onPayload?.(payload as never, model as never);
capturedPayload = payload;
return {} as never;
};
const googleHooks = buildProviderStreamFamilyHooks("google-thinking");
googleHooks.wrapStreamFn?.({
streamFn: baseStreamFn,
thinkingLevel: "high",
} as never)(
{ api: "google-generative-ai", id: "gemini-3.1-pro-preview" } as never,
{} as never,
{},
);
expect(capturedPayload).toMatchObject({
config: { thinkingConfig: { thinkingLevel: "HIGH" } },
});
const googleThinkingConfig = (
(capturedPayload as Record<string, unknown>).config as Record<string, unknown>
).thinkingConfig as Record<string, unknown>;
expect(googleThinkingConfig).not.toHaveProperty("thinkingBudget");
const minimaxHooks = buildProviderStreamFamilyHooks("minimax-fast-mode");
minimaxHooks.wrapStreamFn?.({
streamFn: baseStreamFn,
extraParams: { fastMode: true },
} as never)(
{
api: "anthropic-messages",
provider: "minimax",
id: "MiniMax-M2.7",
} as never,
{} as never,
{},
);
expect(capturedModelId).toBe("MiniMax-M2.7-highspeed");
const moonshotHooks = buildProviderStreamFamilyHooks("moonshot-thinking");
moonshotHooks.wrapStreamFn?.({
streamFn: baseStreamFn,
thinkingLevel: "off",
} as never)(
{ api: "openai-completions", id: "kimi-k2.5" } as never,
{} as never,
{},
);
expect(capturedPayload).toMatchObject({
config: { thinkingConfig: { thinkingBudget: -1 } },
thinking: { type: "disabled" },
});
const toolStreamHooks = buildProviderStreamFamilyHooks("tool-stream-default-on");
toolStreamHooks.wrapStreamFn?.({
streamFn: baseStreamFn,
extraParams: {},
} as never)({ id: "glm-4.7" } as never, {} as never, {});
expect(capturedPayload).toMatchObject({
config: { thinkingConfig: { thinkingBudget: -1 } },
tool_stream: true,
});
toolStreamHooks.wrapStreamFn?.({
streamFn: baseStreamFn,
extraParams: { tool_stream: false },
} as never)({ id: "glm-4.7" } as never, {} as never, {});
expect(capturedPayload).toMatchObject({
config: { thinkingConfig: { thinkingBudget: -1 } },
});
expect(capturedPayload).not.toHaveProperty("tool_stream");
});
});

71
src/plugin-sdk/provider-stream.ts
View File

@@ -1,4 +1,22 @@
import type { StreamFn } from "@mariozechner/pi-agent-core";
import type { ProviderPlugin } from "../plugins/types.js";
import type { ProviderWrapStreamFnContext } from "./plugin-entry.js";
import {
createGoogleThinkingPayloadWrapper,
sanitizeGoogleThinkingPayload,
} from "../agents/pi-embedded-runner/google-stream-wrappers.js";
import { createMinimaxFastModeWrapper } from "../agents/pi-embedded-runner/minimax-stream-wrappers.js";
import {
createMoonshotThinkingWrapper,
resolveMoonshotThinkingType,
} from "../agents/pi-embedded-runner/moonshot-thinking-stream-wrappers.js";
import {
createKilocodeWrapper,
createOpenRouterSystemCacheWrapper,
createOpenRouterWrapper,
isProxyReasoningUnsupported,
} from "../agents/pi-embedded-runner/proxy-stream-wrappers.js";
import { createToolStreamWrapper, createZaiToolStreamWrapper } from "../agents/pi-embedded-runner/zai-stream-wrappers.js";
export type ProviderStreamWrapperFactory =
| ((streamFn: StreamFn | undefined) => StreamFn | undefined)
@@ -16,6 +34,46 @@ export function composeProviderStreamWrappers(
);
}
export type ProviderStreamFamily =
| "google-thinking"
| "moonshot-thinking"
| "minimax-fast-mode"
| "tool-stream-default-on";
type ProviderStreamFamilyHooks = Pick<ProviderPlugin, "wrapStreamFn">;
export function buildProviderStreamFamilyHooks(
family: ProviderStreamFamily,
): ProviderStreamFamilyHooks {
switch (family) {
case "google-thinking":
return {
wrapStreamFn: (ctx: ProviderWrapStreamFnContext) =>
createGoogleThinkingPayloadWrapper(ctx.streamFn, ctx.thinkingLevel),
};
case "moonshot-thinking":
return {
wrapStreamFn: (ctx: ProviderWrapStreamFnContext) => {
const thinkingType = resolveMoonshotThinkingType({
configuredThinking: ctx.extraParams?.thinking,
thinkingLevel: ctx.thinkingLevel,
});
return createMoonshotThinkingWrapper(ctx.streamFn, thinkingType);
},
};
case "minimax-fast-mode":
return {
wrapStreamFn: (ctx: ProviderWrapStreamFnContext) =>
createMinimaxFastModeWrapper(ctx.streamFn, ctx.extraParams?.fastMode === true),
};
case "tool-stream-default-on":
return {
wrapStreamFn: (ctx: ProviderWrapStreamFnContext) =>
createToolStreamWrapper(ctx.streamFn, ctx.extraParams?.tool_stream !== false),
};
}
}
// Public stream-wrapper helpers for provider plugins.
export {
@@ -38,18 +96,14 @@ export {
export {
createGoogleThinkingPayloadWrapper,
sanitizeGoogleThinkingPayload,
} from "../agents/pi-embedded-runner/google-stream-wrappers.js";
export { createMinimaxFastModeWrapper } from "../agents/pi-embedded-runner/minimax-stream-wrappers.js";
export {
createMinimaxFastModeWrapper,
createKilocodeWrapper,
createOpenRouterSystemCacheWrapper,
createOpenRouterWrapper,
isProxyReasoningUnsupported,
} from "../agents/pi-embedded-runner/proxy-stream-wrappers.js";
export {
createMoonshotThinkingWrapper,
resolveMoonshotThinkingType,
} from "../agents/pi-embedded-runner/moonshot-thinking-stream-wrappers.js";
};
export {
createOpenAIAttributionHeadersWrapper,
createCodexNativeWebSearchWrapper,
@@ -64,10 +118,7 @@ export {
resolveOpenAITextVerbosity,
} from "../agents/pi-embedded-runner/openai-stream-wrappers.js";
export { streamWithPayloadPatch } from "../agents/pi-embedded-runner/stream-payload-utils.js";
export {
createToolStreamWrapper,
createZaiToolStreamWrapper,
} from "../agents/pi-embedded-runner/zai-stream-wrappers.js";
export { createToolStreamWrapper, createZaiToolStreamWrapper };
export {
getOpenRouterModelCapabilities,
loadOpenRouterModelCapabilities,