From ac88a39accdf915c412a0f6904425989a24bc885 Mon Sep 17 00:00:00 2001 From: Ayaan Zaidi Date: Tue, 10 Mar 2026 20:29:03 +0530 Subject: [PATCH] fix: align pi-ai 0.57.1 oauth imports and payload hooks --- src/agents/anthropic-payload-log.test.ts | 2 +- src/agents/anthropic-payload-log.ts | 2 +- ...auth.openai-codex-refresh-fallback.test.ts | 18 ++++------ src/agents/auth-profiles/oauth.ts | 2 +- src/agents/openai-ws-stream.ts | 8 +++-- .../pi-embedded-runner-extraparams.test.ts | 34 +++++++++---------- .../anthropic-stream-wrappers.ts | 2 +- .../extra-params.kilocode.test.ts | 8 ++--- ...ra-params.openrouter-cache-control.test.ts | 2 +- src/agents/pi-embedded-runner/extra-params.ts | 6 ++-- .../extra-params.zai-tool-stream.test.ts | 2 +- .../moonshot-stream-wrappers.ts | 4 +-- .../openai-stream-wrappers.ts | 4 +-- .../proxy-stream-wrappers.ts | 6 ++-- .../pi-embedded-runner/run/attempt.test.ts | 2 +- src/agents/pi-embedded-runner/run/attempt.ts | 4 +-- src/commands/openai-codex-oauth.test.ts | 2 +- src/commands/openai-codex-oauth.ts | 2 +- src/tts/tts.test.ts | 3 ++ 19 files changed, 58 insertions(+), 55 deletions(-) diff --git a/src/agents/anthropic-payload-log.test.ts b/src/agents/anthropic-payload-log.test.ts index 037093fbbf5..fb3cf18e47d 100644 --- a/src/agents/anthropic-payload-log.test.ts +++ b/src/agents/anthropic-payload-log.test.ts @@ -29,7 +29,7 @@ describe("createAnthropicPayloadLogger", () => { ], }; const streamFn: StreamFn = ((model, __, options) => { - options?.onPayload?.(payload); + options?.onPayload?.(payload, model); return {} as never; }) as StreamFn; diff --git a/src/agents/anthropic-payload-log.ts b/src/agents/anthropic-payload-log.ts index d80ed551179..2eb5d62e770 100644 --- a/src/agents/anthropic-payload-log.ts +++ b/src/agents/anthropic-payload-log.ts @@ -145,7 +145,7 @@ export function createAnthropicPayloadLogger(params: { payload: redactedPayload, payloadDigest: digest(redactedPayload), }); - return options?.onPayload?.(payload); + return options?.onPayload?.(payload, model); }; return streamFn(model, context, { ...options, diff --git a/src/agents/auth-profiles/oauth.openai-codex-refresh-fallback.test.ts b/src/agents/auth-profiles/oauth.openai-codex-refresh-fallback.test.ts index 9d47be8c79e..23381d89a05 100644 --- a/src/agents/auth-profiles/oauth.openai-codex-refresh-fallback.test.ts +++ b/src/agents/auth-profiles/oauth.openai-codex-refresh-fallback.test.ts @@ -17,17 +17,13 @@ const { getOAuthApiKeyMock } = vi.hoisted(() => ({ }), })); -vi.mock("@mariozechner/pi-ai", async () => { - const actual = await vi.importActual("@mariozechner/pi-ai"); - return { - ...actual, - getOAuthApiKey: getOAuthApiKeyMock, - getOAuthProviders: () => [ - { id: "openai-codex", envApiKey: "OPENAI_API_KEY", oauthTokenEnv: "OPENAI_OAUTH_TOKEN" }, // pragma: allowlist secret - { id: "anthropic", envApiKey: "ANTHROPIC_API_KEY", oauthTokenEnv: "ANTHROPIC_OAUTH_TOKEN" }, // pragma: allowlist secret - ], - }; -}); +vi.mock("@mariozechner/pi-ai/oauth", () => ({ + getOAuthApiKey: getOAuthApiKeyMock, + getOAuthProviders: () => [ + { id: "openai-codex", envApiKey: "OPENAI_API_KEY", oauthTokenEnv: "OPENAI_OAUTH_TOKEN" }, // pragma: allowlist secret + { id: "anthropic", envApiKey: "ANTHROPIC_API_KEY", oauthTokenEnv: "ANTHROPIC_OAUTH_TOKEN" }, // pragma: allowlist secret + ], +})); function createExpiredOauthStore(params: { profileId: string; diff --git a/src/agents/auth-profiles/oauth.ts b/src/agents/auth-profiles/oauth.ts index a67e8e6a6bb..072b3a77246 100644 --- a/src/agents/auth-profiles/oauth.ts +++ b/src/agents/auth-profiles/oauth.ts @@ -1,5 +1,5 @@ import type { OAuthCredentials, OAuthProvider } from "@mariozechner/pi-ai"; -import { getOAuthApiKey, getOAuthProviders } from "@mariozechner/pi-ai"; +import { getOAuthApiKey, getOAuthProviders } from "@mariozechner/pi-ai/oauth"; import { loadConfig, type OpenClawConfig } from "../../config/config.js"; import { coerceSecretRef } from "../../config/types.secrets.js"; import { withFileLock } from "../../infra/file-lock.js"; diff --git a/src/agents/openai-ws-stream.ts b/src/agents/openai-ws-stream.ts index 9228fd92d46..dd82ced9e95 100644 --- a/src/agents/openai-ws-stream.ts +++ b/src/agents/openai-ws-stream.ts @@ -604,10 +604,14 @@ export function createOpenAIWebSocketStreamFn( ...(prevResponseId ? { previous_response_id: prevResponseId } : {}), ...extraParams, }; - options?.onPayload?.(payload); + const nextPayload = await options?.onPayload?.(payload, model); + const requestPayload = + nextPayload && typeof nextPayload === "object" + ? (nextPayload as Parameters[0]) + : (payload as Parameters[0]); try { - session.manager.send(payload as Parameters[0]); + session.manager.send(requestPayload); } catch (sendErr) { if (transport === "websocket") { throw sendErr instanceof Error ? sendErr : new Error(String(sendErr)); diff --git a/src/agents/pi-embedded-runner-extraparams.test.ts b/src/agents/pi-embedded-runner-extraparams.test.ts index 6689b3426cf..232cdfcaa0b 100644 --- a/src/agents/pi-embedded-runner-extraparams.test.ts +++ b/src/agents/pi-embedded-runner-extraparams.test.ts @@ -208,7 +208,7 @@ describe("applyExtraParamsToAgent", () => { }) { const payload = params.payload ?? { store: false }; const baseStreamFn: StreamFn = (model, _context, options) => { - options?.onPayload?.(payload); + options?.onPayload?.(payload, model); return {} as ReturnType; }; const agent = { streamFn: baseStreamFn }; @@ -233,7 +233,7 @@ describe("applyExtraParamsToAgent", () => { }) { const payload = params.payload ?? {}; const baseStreamFn: StreamFn = (model, _context, options) => { - options?.onPayload?.(payload); + options?.onPayload?.(payload, model); return {} as ReturnType; }; const agent = { streamFn: baseStreamFn }; @@ -276,7 +276,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { model: "deepseek/deepseek-r1" }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -308,7 +308,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = {}; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -332,7 +332,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { reasoning_effort: "high" }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -357,7 +357,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { reasoning: { max_tokens: 256 } }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -381,7 +381,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { reasoning_effort: "medium" }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -588,7 +588,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { thinking: "off" }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -619,7 +619,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { thinking: "off" }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -650,7 +650,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = {}; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -674,7 +674,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { tool_choice: "required" }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -699,7 +699,7 @@ describe("applyExtraParamsToAgent", () => { const payloads: Record[] = []; const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = {}; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -749,7 +749,7 @@ describe("applyExtraParamsToAgent", () => { ], tool_choice: { type: "tool", name: "read" }, }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -793,7 +793,7 @@ describe("applyExtraParamsToAgent", () => { }, ], }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -832,7 +832,7 @@ describe("applyExtraParamsToAgent", () => { }, ], }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -896,7 +896,7 @@ describe("applyExtraParamsToAgent", () => { }, }, }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; @@ -943,7 +943,7 @@ describe("applyExtraParamsToAgent", () => { }, }, }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloads.push(payload); return {} as ReturnType; }; diff --git a/src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts b/src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts index 66718b9e0aa..df43d2570c7 100644 --- a/src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts +++ b/src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts @@ -298,7 +298,7 @@ export function createAnthropicToolPayloadCompatibilityWrapper( ); } } - return originalOnPayload?.(payload); + return originalOnPayload?.(payload, model); }, }); }; diff --git a/src/agents/pi-embedded-runner/extra-params.kilocode.test.ts b/src/agents/pi-embedded-runner/extra-params.kilocode.test.ts index 509cdb5edf4..0e2fd5ce93b 100644 --- a/src/agents/pi-embedded-runner/extra-params.kilocode.test.ts +++ b/src/agents/pi-embedded-runner/extra-params.kilocode.test.ts @@ -19,7 +19,7 @@ function applyAndCapture(params: { const baseStreamFn: StreamFn = (_model, _context, options) => { captured.headers = options?.headers; - options?.onPayload?.({}); + options?.onPayload?.({}, _model); return createAssistantMessageEventStream(); }; const agent = { streamFn: baseStreamFn }; @@ -97,7 +97,7 @@ describe("extra-params: Kilocode kilo/auto reasoning", () => { const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { reasoning_effort: "high" }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); capturedPayload = payload; return createAssistantMessageEventStream(); }; @@ -125,7 +125,7 @@ describe("extra-params: Kilocode kilo/auto reasoning", () => { const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = {}; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); capturedPayload = payload; return createAssistantMessageEventStream(); }; @@ -158,7 +158,7 @@ describe("extra-params: Kilocode kilo/auto reasoning", () => { const baseStreamFn: StreamFn = (_model, _context, options) => { const payload: Record = { reasoning_effort: "high" }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); capturedPayload = payload; return createAssistantMessageEventStream(); }; diff --git a/src/agents/pi-embedded-runner/extra-params.openrouter-cache-control.test.ts b/src/agents/pi-embedded-runner/extra-params.openrouter-cache-control.test.ts index 71af916ccac..58af2239a3d 100644 --- a/src/agents/pi-embedded-runner/extra-params.openrouter-cache-control.test.ts +++ b/src/agents/pi-embedded-runner/extra-params.openrouter-cache-control.test.ts @@ -13,7 +13,7 @@ type StreamPayload = { function runOpenRouterPayload(payload: StreamPayload, modelId: string) { const baseStreamFn: StreamFn = (_model, _context, options) => { - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); return createAssistantMessageEventStream(); }; const agent = { streamFn: baseStreamFn }; diff --git a/src/agents/pi-embedded-runner/extra-params.ts b/src/agents/pi-embedded-runner/extra-params.ts index 6e261463d4a..8f36792f393 100644 --- a/src/agents/pi-embedded-runner/extra-params.ts +++ b/src/agents/pi-embedded-runner/extra-params.ts @@ -230,7 +230,7 @@ function createGoogleThinkingPayloadWrapper( thinkingLevel, }); } - return onPayload?.(payload); + return onPayload?.(payload, model); }, }); }; @@ -263,7 +263,7 @@ function createZaiToolStreamWrapper( // Inject tool_stream: true for Z.AI API (payload as Record).tool_stream = true; } - return originalOnPayload?.(payload); + return originalOnPayload?.(payload, model); }, }); }; @@ -310,7 +310,7 @@ function createParallelToolCallsWrapper( if (payload && typeof payload === "object") { (payload as Record).parallel_tool_calls = enabled; } - return originalOnPayload?.(payload); + return originalOnPayload?.(payload, model); }, }); }; 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 index 2dab69cd15a..f7262a66798 100644 --- 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 @@ -22,7 +22,7 @@ type ToolStreamCase = { function runToolStreamCase(params: ToolStreamCase) { const payload: Record = { model: params.model.id, messages: [] }; const baseStreamFn: StreamFn = (model, _context, options) => { - options?.onPayload?.(payload); + options?.onPayload?.(payload, model); return {} as ReturnType; }; const agent = { streamFn: baseStreamFn }; diff --git a/src/agents/pi-embedded-runner/moonshot-stream-wrappers.ts b/src/agents/pi-embedded-runner/moonshot-stream-wrappers.ts index aa43260e55e..282b0960a9d 100644 --- a/src/agents/pi-embedded-runner/moonshot-stream-wrappers.ts +++ b/src/agents/pi-embedded-runner/moonshot-stream-wrappers.ts @@ -60,7 +60,7 @@ export function createSiliconFlowThinkingWrapper(baseStreamFn: StreamFn | undefi payloadObj.thinking = null; } } - return originalOnPayload?.(payload); + return originalOnPayload?.(payload, model); }, }); }; @@ -106,7 +106,7 @@ export function createMoonshotThinkingWrapper( payloadObj.tool_choice = "auto"; } } - return originalOnPayload?.(payload); + return originalOnPayload?.(payload, model); }, }); }; diff --git a/src/agents/pi-embedded-runner/openai-stream-wrappers.ts b/src/agents/pi-embedded-runner/openai-stream-wrappers.ts index 7c54a7dbc37..3fc46dac0ae 100644 --- a/src/agents/pi-embedded-runner/openai-stream-wrappers.ts +++ b/src/agents/pi-embedded-runner/openai-stream-wrappers.ts @@ -197,7 +197,7 @@ export function createOpenAIResponsesContextManagementWrapper( compactThreshold, }); } - return originalOnPayload?.(payload); + return originalOnPayload?.(payload, model); }, }); }; @@ -226,7 +226,7 @@ export function createOpenAIServiceTierWrapper( payloadObj.service_tier = serviceTier; } } - return originalOnPayload?.(payload); + return originalOnPayload?.(payload, model); }, }); }; diff --git a/src/agents/pi-embedded-runner/proxy-stream-wrappers.ts b/src/agents/pi-embedded-runner/proxy-stream-wrappers.ts index a5f9f5b1d85..4f77c31cfdd 100644 --- a/src/agents/pi-embedded-runner/proxy-stream-wrappers.ts +++ b/src/agents/pi-embedded-runner/proxy-stream-wrappers.ts @@ -92,7 +92,7 @@ export function createOpenRouterSystemCacheWrapper(baseStreamFn: StreamFn | unde } } } - return originalOnPayload?.(payload); + return originalOnPayload?.(payload, model); }, }); }; @@ -113,7 +113,7 @@ export function createOpenRouterWrapper( }, onPayload: (payload) => { normalizeProxyReasoningPayload(payload, thinkingLevel); - return onPayload?.(payload); + return onPayload?.(payload, model); }, }); }; @@ -138,7 +138,7 @@ export function createKilocodeWrapper( }, onPayload: (payload) => { normalizeProxyReasoningPayload(payload, thinkingLevel); - return onPayload?.(payload); + return onPayload?.(payload, model); }, }); }; diff --git a/src/agents/pi-embedded-runner/run/attempt.test.ts b/src/agents/pi-embedded-runner/run/attempt.test.ts index 70bd3242f7c..9821adc0e0b 100644 --- a/src/agents/pi-embedded-runner/run/attempt.test.ts +++ b/src/agents/pi-embedded-runner/run/attempt.test.ts @@ -520,7 +520,7 @@ describe("wrapOllamaCompatNumCtx", () => { let payloadSeen: Record | undefined; const baseFn = vi.fn((_model, _context, options) => { const payload: Record = { options: { temperature: 0.1 } }; - options?.onPayload?.(payload); + options?.onPayload?.(payload, _model); payloadSeen = payload; return {} as never; }); diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 80c674ae7c1..084a6d39746 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -230,14 +230,14 @@ export function wrapOllamaCompatNumCtx(baseFn: StreamFn | undefined, numCtx: num ...options, onPayload: (payload: unknown) => { if (!payload || typeof payload !== "object") { - return options?.onPayload?.(payload); + return options?.onPayload?.(payload, model); } const payloadRecord = payload as Record; if (!payloadRecord.options || typeof payloadRecord.options !== "object") { payloadRecord.options = {}; } (payloadRecord.options as Record).num_ctx = numCtx; - return options?.onPayload?.(payload); + return options?.onPayload?.(payload, model); }, }); } diff --git a/src/commands/openai-codex-oauth.test.ts b/src/commands/openai-codex-oauth.test.ts index abe71d0bd42..43f1ac41f8a 100644 --- a/src/commands/openai-codex-oauth.test.ts +++ b/src/commands/openai-codex-oauth.test.ts @@ -9,7 +9,7 @@ const mocks = vi.hoisted(() => ({ formatOpenAIOAuthTlsPreflightFix: vi.fn(), })); -vi.mock("@mariozechner/pi-ai", () => ({ +vi.mock("@mariozechner/pi-ai/oauth", () => ({ loginOpenAICodex: mocks.loginOpenAICodex, })); diff --git a/src/commands/openai-codex-oauth.ts b/src/commands/openai-codex-oauth.ts index 72a13f654cf..1f6a8f9cde8 100644 --- a/src/commands/openai-codex-oauth.ts +++ b/src/commands/openai-codex-oauth.ts @@ -1,5 +1,5 @@ import type { OAuthCredentials } from "@mariozechner/pi-ai"; -import { loginOpenAICodex } from "@mariozechner/pi-ai"; +import { loginOpenAICodex } from "@mariozechner/pi-ai/oauth"; import type { RuntimeEnv } from "../runtime.js"; import type { WizardPrompter } from "../wizard/prompts.js"; import { createVpsAwareOAuthHandlers } from "./oauth-flow.js"; diff --git a/src/tts/tts.test.ts b/src/tts/tts.test.ts index 733d34f5757..f3b5d8ce0ee 100644 --- a/src/tts/tts.test.ts +++ b/src/tts/tts.test.ts @@ -9,6 +9,9 @@ import * as tts from "./tts.js"; vi.mock("@mariozechner/pi-ai", () => ({ completeSimple: vi.fn(), +})); + +vi.mock("@mariozechner/pi-ai/oauth", () => ({ // Some auth helpers import oauth provider metadata at module load time. getOAuthProviders: () => [], getOAuthApiKey: vi.fn(async () => null),