From edbc68e9f1e4652c9c81b0333ad624e9c150d867 Mon Sep 17 00:00:00 2001 From: tian Xiao Date: Tue, 17 Feb 2026 00:30:33 +0900 Subject: [PATCH] feat: support Z.AI tool_stream for real-time tool call streaming Add support for Z.AI's native tool_stream parameter to enable real-time visibility into model reasoning and tool call execution. - Automatically inject tool_stream=true for zai/z-ai providers - Allow disabling via params.tool_stream: false in model config - Follows existing pattern of OpenRouter and OpenAI wrappers This enables Z.AI API features described in: https://docs.z.ai/api-reference#streaming AI-assisted: Claude (OpenClaw agent) helped write this implementation. Testing: lightly tested (code review + pattern matching existing wrappers) Closes #18135 --- src/agents/pi-embedded-runner/extra-params.ts | 43 +++++++++ .../extra-params.zai-tool-stream.test.ts | 94 +++++++++++++++++++ 2 files changed, 137 insertions(+) create mode 100644 src/agents/pi-embedded-runner/extra-params.zai-tool-stream.test.ts diff --git a/src/agents/pi-embedded-runner/extra-params.ts b/src/agents/pi-embedded-runner/extra-params.ts index 08cef5491ba..70154e5b550 100644 --- a/src/agents/pi-embedded-runner/extra-params.ts +++ b/src/agents/pi-embedded-runner/extra-params.ts @@ -172,6 +172,39 @@ function createOpenRouterHeadersWrapper(baseStreamFn: StreamFn | undefined): Str }); } +/** + * Create a streamFn wrapper that injects tool_stream=true for Z.AI providers. + * + * Z.AI's API supports the `tool_stream` parameter to enable real-time streaming + * of tool call arguments and reasoning content. When enabled, the API returns + * progressive tool_call deltas, allowing users to see tool execution in real-time. + * + * @see https://docs.z.ai/api-reference#streaming + */ +function createZaiToolStreamWrapper( + baseStreamFn: StreamFn | undefined, + enabled: boolean, +): StreamFn { + const underlying = baseStreamFn ?? streamSimple; + return (model, context, options) => { + if (!enabled) { + return underlying(model, context, options); + } + + const originalOnPayload = options?.onPayload; + return underlying(model, context, { + ...options, + onPayload: (payload) => { + if (payload && typeof payload === "object") { + // Inject tool_stream: true for Z.AI API + (payload as Record).tool_stream = true; + } + originalOnPayload?.(payload); + }, + }); + }; +} + /** * Apply extra params (like temperature) to an agent's streamFn. * Also adds OpenRouter app attribution headers when using the OpenRouter provider. @@ -209,6 +242,16 @@ export function applyExtraParamsToAgent( agent.streamFn = createOpenRouterHeadersWrapper(agent.streamFn); } + // Enable Z.AI tool_stream for real-time tool call streaming. + // Enabled by default for Z.AI provider, can be disabled via params.tool_stream: false + if (provider === "zai" || provider === "z-ai") { + const toolStreamEnabled = merged?.tool_stream !== false; + if (toolStreamEnabled) { + log.debug(`enabling Z.AI tool_stream for ${provider}/${modelId}`); + agent.streamFn = createZaiToolStreamWrapper(agent.streamFn, true); + } + } + // Work around upstream pi-ai hardcoding `store: false` for Responses API. // Force `store=true` for direct OpenAI/OpenAI Codex providers so multi-turn // server-side conversation state is preserved. diff --git a/src/agents/pi-embedded-runner/extra-params.zai-tool-stream.test.ts b/src/agents/pi-embedded-runner/extra-params.zai-tool-stream.test.ts new file mode 100644 index 00000000000..d1903ef9a80 --- /dev/null +++ b/src/agents/pi-embedded-runner/extra-params.zai-tool-stream.test.ts @@ -0,0 +1,94 @@ +import { describe, expect, it, vi } from "vitest"; +import type { StreamFn } from "@mariozechner/pi-agent-core"; +import type { Context, Model } from "@mariozechner/pi-ai"; +import { applyExtraParamsToAgent } from "./extra-params.js"; + +// Mock streamSimple for testing +vi.mock("@mariozechner/pi-ai", () => ({ + streamSimple: vi.fn(() => ({ + push: vi.fn(), + result: vi.fn(), + })), +})); + +describe("extra-params: Z.AI tool_stream support", () => { + it("should inject tool_stream=true for zai provider by default", () => { + const capturedPayloads: unknown[] = []; + const mockStreamFn: StreamFn = vi.fn((model, context, options) => { + // Capture the payload that would be sent + options?.onPayload?.({ model: model.id, messages: [] }); + return { + push: vi.fn(), + result: vi.fn().mockResolvedValue({ + role: "assistant", + content: [{ type: "text", text: "ok" }], + stopReason: "stop", + }), + } as any; + }); + + const agent = { streamFn: mockStreamFn }; + const cfg = { + agents: { + defaults: {}, + }, + }; + + applyExtraParamsToAgent(agent, cfg as any, "zai", "glm-5"); + + // The streamFn should be wrapped + expect(agent.streamFn).toBeDefined(); + expect(agent.streamFn).not.toBe(mockStreamFn); + }); + + it("should not inject tool_stream for non-zai providers", () => { + const mockStreamFn: StreamFn = vi.fn(() => ({ + push: vi.fn(), + result: vi.fn().mockResolvedValue({ + role: "assistant", + content: [{ type: "text", text: "ok" }], + stopReason: "stop", + }), + } as any)); + + const agent = { streamFn: mockStreamFn }; + const cfg = {}; + + applyExtraParamsToAgent(agent, cfg as any, "anthropic", "claude-opus-4-6"); + + // Should remain unchanged (except for OpenAI wrapper) + expect(agent.streamFn).toBeDefined(); + }); + + it("should allow disabling tool_stream via params", () => { + const mockStreamFn: StreamFn = vi.fn(() => ({ + push: vi.fn(), + result: vi.fn().mockResolvedValue({ + role: "assistant", + content: [{ type: "text", text: "ok" }], + stopReason: "stop", + }), + } as any)); + + const agent = { streamFn: mockStreamFn }; + const cfg = { + agents: { + defaults: { + models: { + "zai/glm-5": { + params: { + tool_stream: false, + }, + }, + }, + }, + }, + }; + + applyExtraParamsToAgent(agent, cfg as any, "zai", "glm-5"); + + // The tool_stream wrapper should be applied but with enabled=false + // In this case, it should just return the underlying streamFn + expect(agent.streamFn).toBeDefined(); + }); +});