fix: align pi-ai 0.57.1 oauth imports and payload hooks

2026-03-12 07:20:45 +00:00 · 2026-03-10 20:29:03 +05:30
parent f50fc2966b
commit ac88a39acc
19 changed files with 58 additions and 55 deletions
--- 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;

--- 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,
--- 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<typeof import("@mariozechner/pi-ai")>("@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;
--- 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";
--- 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<OpenAIWebSocketManager["send"]>[0])
+          : (payload as Parameters<OpenAIWebSocketManager["send"]>[0]);

      try {
-        session.manager.send(payload as Parameters<OpenAIWebSocketManager["send"]>[0]);
+        session.manager.send(requestPayload);
      } catch (sendErr) {
        if (transport === "websocket") {
          throw sendErr instanceof Error ? sendErr : new Error(String(sendErr));
--- 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<StreamFn>;
    };
    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<StreamFn>;
    };
    const agent = { streamFn: baseStreamFn };
@@ -276,7 +276,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = { model: "deepseek/deepseek-r1" };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -308,7 +308,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = {};
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -332,7 +332,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = { reasoning_effort: "high" };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -357,7 +357,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = { reasoning: { max_tokens: 256 } };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -381,7 +381,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = { reasoning_effort: "medium" };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -588,7 +588,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = { thinking: "off" };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -619,7 +619,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = { thinking: "off" };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -650,7 +650,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = {};
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -674,7 +674,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = { tool_choice: "required" };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -699,7 +699,7 @@ describe("applyExtraParamsToAgent", () => {
    const payloads: Record<string, unknown>[] = [];
    const baseStreamFn: StreamFn = (_model, _context, options) => {
      const payload: Record<string, unknown> = {};
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -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<StreamFn>;
    };
@@ -793,7 +793,7 @@ describe("applyExtraParamsToAgent", () => {
          },
        ],
      };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -832,7 +832,7 @@ describe("applyExtraParamsToAgent", () => {
          },
        ],
      };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -896,7 +896,7 @@ describe("applyExtraParamsToAgent", () => {
          },
        },
      };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
@@ -943,7 +943,7 @@ describe("applyExtraParamsToAgent", () => {
          },
        },
      };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloads.push(payload);
      return {} as ReturnType<StreamFn>;
    };
--- 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);
      },
    });
  };
--- 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<string, unknown> = { 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<string, unknown> = {};
-      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<string, unknown> = { reasoning_effort: "high" };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      capturedPayload = payload;
      return createAssistantMessageEventStream();
    };
--- 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 };
--- 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<string, unknown>).tool_stream = true;
        }
-        return originalOnPayload?.(payload);
+        return originalOnPayload?.(payload, model);
      },
    });
  };
@@ -310,7 +310,7 @@ function createParallelToolCallsWrapper(
        if (payload && typeof payload === "object") {
          (payload as Record<string, unknown>).parallel_tool_calls = enabled;
        }
-        return originalOnPayload?.(payload);
+        return originalOnPayload?.(payload, model);
      },
    });
  };
--- 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<string, unknown> = { model: params.model.id, messages: [] };
  const baseStreamFn: StreamFn = (model, _context, options) => {
-    options?.onPayload?.(payload);
+    options?.onPayload?.(payload, model);
    return {} as ReturnType<StreamFn>;
  };
  const agent = { streamFn: baseStreamFn };
--- 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);
      },
    });
  };
--- 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);
      },
    });
  };
--- 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);
      },
    });
  };
--- 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<string, unknown> | undefined;
    const baseFn = vi.fn((_model, _context, options) => {
      const payload: Record<string, unknown> = { options: { temperature: 0.1 } };
-      options?.onPayload?.(payload);
+      options?.onPayload?.(payload, _model);
      payloadSeen = payload;
      return {} as never;
    });
--- 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<string, unknown>;
        if (!payloadRecord.options || typeof payloadRecord.options !== "object") {
          payloadRecord.options = {};
        }
        (payloadRecord.options as Record<string, unknown>).num_ctx = numCtx;
-        return options?.onPayload?.(payload);
+        return options?.onPayload?.(payload, model);
      },
    });
 }
--- 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,
 }));

--- 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";
--- 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),